Skip to content

Commit 85ebb89

Browse files
pwilkino7si
pwilkin
authored and
o7si
committed
grammar: optimize big grammar matching by using left-factorization
1 parent 8139cba commit 85ebb89

3 files changed

Lines changed: 697 additions & 18 deletions

File tree

‎src/llama-grammar.cpp‎

Lines changed: 209 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -948,6 +948,188 @@ static void llama_grammar_advance_stack(
948948
}
949949
}
950950

951+
// Left-factor grammar rules to reduce redundant parallel stacks.
952+
//
953+
// When a rule has many alternates (e.g. tool-choice ::= tc0 | tc1 | ... | tc119)
954+
// and the referenced rules share a common element-content prefix, extract that
955+
// prefix so the parser processes it once instead of N times.
956+
//
957+
// Handles two cases:
958+
// 1. Inline common prefix: rule = "ab1" | "ab2" → rule = "ab" rule', rule' = "1" | "2"
959+
// 2. Cross-rule prefix: rule = A | B | C where A,B,C start with same elements
960+
// → rule = common_prefix rule', rule' = A_suffix | B_suffix | C_suffix
961+
static void llama_grammar_left_factor(llama_grammar_rules & rules) {
962+
auto elem_eq = [](const llama_grammar_element & a, const llama_grammar_element & b) {
963+
return a.type == b.type && a.value == b.value;
964+
};
965+
966+
struct alt_info {
967+
size_t rule_idx;
968+
size_t start;
969+
size_t end;
970+
};
971+
972+
bool changed = true;
973+
while (changed) {
974+
changed = false;
975+
const size_t n_rules = rules.size();
976+
977+
for (size_t ri = 0; ri < n_rules; ri++) {
978+
const auto & rule = rules[ri];
979+
980+
// Split rule into alternates
981+
std::vector<std::pair<size_t, size_t>> alt_ranges;
982+
size_t alt_start = 0;
983+
for (size_t j = 0; j < rule.size(); j++) {
984+
if (rule[j].type == LLAMA_GRETYPE_ALT || rule[j].type == LLAMA_GRETYPE_END) {
985+
alt_ranges.push_back({alt_start, j});
986+
alt_start = j + 1;
987+
}
988+
}
989+
990+
if (alt_ranges.size() < 4) {
991+
continue;
992+
}
993+
994+
// Resolve each alternate: if it's a single RULE_REF, follow it to get the
995+
// actual element sequence from the referenced rule.
996+
std::vector<alt_info> resolved;
997+
resolved.reserve(alt_ranges.size());
998+
bool all_single_refs = true;
999+
1000+
for (const auto & [astart, aend] : alt_ranges) {
1001+
size_t len = aend - astart;
1002+
if (len == 1 && rule[astart].type == LLAMA_GRETYPE_RULE_REF) {
1003+
size_t ref_id = static_cast<size_t>(rule[astart].value);
1004+
if (ref_id < rules.size()) {
1005+
size_t ref_end = 0;
1006+
for (size_t k = 0; k < rules[ref_id].size(); k++) {
1007+
if (rules[ref_id][k].type == LLAMA_GRETYPE_ALT ||
1008+
rules[ref_id][k].type == LLAMA_GRETYPE_END) {
1009+
ref_end = k;
1010+
break;
1011+
}
1012+
}
1013+
bool single_alt = (rules[ref_id][ref_end].type == LLAMA_GRETYPE_END);
1014+
if (single_alt) {
1015+
resolved.push_back({ref_id, 0, ref_end});
1016+
continue;
1017+
}
1018+
}
1019+
}
1020+
all_single_refs = false;
1021+
resolved.push_back({ri, astart, aend});
1022+
}
1023+
1024+
if (!all_single_refs) {
1025+
// Inline-only factoring
1026+
size_t prefix_len = 0;
1027+
bool prefix_match = true;
1028+
while (prefix_match) {
1029+
const auto & first = alt_ranges[0];
1030+
if (first.first + prefix_len >= first.second) {
1031+
break;
1032+
}
1033+
const auto & ref_elem = rule[first.first + prefix_len];
1034+
for (size_t ai = 1; ai < alt_ranges.size(); ai++) {
1035+
const auto & alt = alt_ranges[ai];
1036+
if (alt.first + prefix_len >= alt.second ||
1037+
!elem_eq(rule[alt.first + prefix_len], ref_elem)) {
1038+
prefix_match = false;
1039+
break;
1040+
}
1041+
}
1042+
if (prefix_match) {
1043+
prefix_len++;
1044+
}
1045+
}
1046+
1047+
if (prefix_len == 0) {
1048+
continue;
1049+
}
1050+
1051+
uint32_t suffix_rule_id = static_cast<uint32_t>(rules.size());
1052+
llama_grammar_rule suffix_rule;
1053+
for (size_t ai = 0; ai < alt_ranges.size(); ai++) {
1054+
if (ai > 0) {
1055+
suffix_rule.push_back({LLAMA_GRETYPE_ALT, 0});
1056+
}
1057+
const auto & alt = alt_ranges[ai];
1058+
for (size_t j = alt.first + prefix_len; j < alt.second; j++) {
1059+
suffix_rule.push_back(rule[j]);
1060+
}
1061+
}
1062+
suffix_rule.push_back({LLAMA_GRETYPE_END, 0});
1063+
1064+
llama_grammar_rule new_rule;
1065+
for (size_t j = 0; j < prefix_len; j++) {
1066+
new_rule.push_back(rule[alt_ranges[0].first + j]);
1067+
}
1068+
new_rule.push_back({LLAMA_GRETYPE_RULE_REF, suffix_rule_id});
1069+
new_rule.push_back({LLAMA_GRETYPE_END, 0});
1070+
1071+
rules.push_back(std::move(suffix_rule));
1072+
rules[ri] = std::move(new_rule);
1073+
changed = true;
1074+
continue;
1075+
}
1076+
1077+
// Cross-rule factoring: all alternates are single RULE_REFs to single-alternate rules.
1078+
size_t prefix_len = 0;
1079+
bool prefix_match = true;
1080+
while (prefix_match) {
1081+
const auto & first = resolved[0];
1082+
if (first.start + prefix_len >= first.end) {
1083+
break;
1084+
}
1085+
const auto & ref_elem = rules[first.rule_idx][first.start + prefix_len];
1086+
1087+
for (size_t ai = 1; ai < resolved.size(); ai++) {
1088+
const auto & alt = resolved[ai];
1089+
if (alt.start + prefix_len >= alt.end ||
1090+
!elem_eq(rules[alt.rule_idx][alt.start + prefix_len], ref_elem)) {
1091+
prefix_match = false;
1092+
break;
1093+
}
1094+
}
1095+
if (prefix_match) {
1096+
prefix_len++;
1097+
}
1098+
}
1099+
1100+
if (prefix_len == 0) {
1101+
continue;
1102+
}
1103+
1104+
uint32_t suffix_choice_id = static_cast<uint32_t>(rules.size());
1105+
llama_grammar_rule suffix_choice;
1106+
1107+
for (size_t ai = 0; ai < resolved.size(); ai++) {
1108+
if (ai > 0) {
1109+
suffix_choice.push_back({LLAMA_GRETYPE_ALT, 0});
1110+
}
1111+
const auto & alt = resolved[ai];
1112+
for (size_t j = alt.start + prefix_len; j < alt.end; j++) {
1113+
suffix_choice.push_back(rules[alt.rule_idx][j]);
1114+
}
1115+
}
1116+
suffix_choice.push_back({LLAMA_GRETYPE_END, 0});
1117+
1118+
llama_grammar_rule new_rule;
1119+
const auto & first = resolved[0];
1120+
for (size_t j = 0; j < prefix_len; j++) {
1121+
new_rule.push_back(rules[first.rule_idx][first.start + j]);
1122+
}
1123+
new_rule.push_back({LLAMA_GRETYPE_RULE_REF, suffix_choice_id});
1124+
new_rule.push_back({LLAMA_GRETYPE_END, 0});
1125+
1126+
rules.push_back(std::move(suffix_choice));
1127+
rules[ri] = std::move(new_rule);
1128+
changed = true;
1129+
}
1130+
}
1131+
}
1132+
9511133
static bool llama_grammar_detect_left_recursion(
9521134
const llama_grammar_rules & rules,
9531135
size_t rule_index,
@@ -1221,12 +1403,16 @@ struct llama_grammar * llama_grammar_init_impl(
12211403
}
12221404
}
12231405

1224-
// Check for left recursion
1225-
std::vector<bool> rules_visited(n_rules);
1226-
std::vector<bool> rules_in_progress(n_rules);
1227-
std::vector<bool> rules_may_be_empty(n_rules);
1406+
// Left-factor rules to reduce redundant parallel stacks
1407+
llama_grammar_left_factor(vec_rules);
1408+
1409+
// Check for left recursion (use post-factoring rule count)
1410+
const size_t n_rules_total = vec_rules.size();
1411+
std::vector<bool> rules_visited(n_rules_total);
1412+
std::vector<bool> rules_in_progress(n_rules_total);
1413+
std::vector<bool> rules_may_be_empty(n_rules_total);
12281414
try {
1229-
for (size_t i = 0; i < n_rules; i++) {
1415+
for (size_t i = 0; i < n_rules_total; i++) {
12301416
if (rules_visited[i]) {
12311417
continue;
12321418
}
@@ -1319,12 +1505,16 @@ struct llama_grammar * llama_grammar_init_impl(
13191505
vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
13201506
}
13211507

1322-
// Check for left recursion
1323-
std::vector<bool> rules_visited(n_rules);
1324-
std::vector<bool> rules_in_progress(n_rules);
1325-
std::vector<bool> rules_may_be_empty(n_rules);
1508+
// Left-factor rules to reduce redundant parallel stacks
1509+
llama_grammar_left_factor(vec_rules);
1510+
1511+
// Check for left recursion (use post-factoring rule count)
1512+
const size_t n_rules_total = vec_rules.size();
1513+
std::vector<bool> rules_visited(n_rules_total);
1514+
std::vector<bool> rules_in_progress(n_rules_total);
1515+
std::vector<bool> rules_may_be_empty(n_rules_total);
13261516
try {
1327-
for (size_t i = 0; i < n_rules; i++) {
1517+
for (size_t i = 0; i < n_rules_total; i++) {
13281518
if (rules_visited[i]) {
13291519
continue;
13301520
}
@@ -1412,14 +1602,15 @@ struct llama_grammar * llama_grammar_clone_impl(const struct llama_grammar & gra
14121602
grammar.trigger_patterns,
14131603
};
14141604

1415-
// redirect elements in stacks to point to new rules
1416-
for (size_t is = 0; is < result->stacks.size(); is++) {
1417-
for (size_t ie = 0; ie < result->stacks[is].size(); ie++) {
1418-
for (size_t ir0 = 0; ir0 < grammar.rules.size(); ir0++) {
1419-
for (size_t ir1 = 0; ir1 < grammar.rules[ir0].size(); ir1++) {
1420-
if (grammar.stacks[is][ie] == &grammar.rules[ir0][ir1]) {
1421-
result->stacks[is][ie] = &result->rules[ir0][ir1];
1422-
}
1605+
// redirect elements in stacks to point to new rules using per-rule pointer offsets
1606+
for (auto & stack : result->stacks) {
1607+
for (auto & elem : stack) {
1608+
for (size_t ir = 0; ir < grammar.rules.size(); ir++) {
1609+
const auto & old_rule = grammar.rules[ir];
1610+
if (!old_rule.empty() && elem >= old_rule.data() && elem < old_rule.data() + old_rule.size()) {
1611+
size_t offset = static_cast<size_t>(elem - old_rule.data());
1612+
elem = result->rules[ir].data() + offset;
1613+
break;
14231614
}
14241615
}
14251616
}

‎tests/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
189189
llama_build_and_test(test-grammar-parser.cpp LABEL behavior LINK llama)
190190
llama_build_and_test(test-grammar-integration.cpp LABEL behavior)
191191
llama_build_and_test(test-llama-grammar.cpp LABEL behavior LINK llama)
192+
llama_build(test-grammar-perf.cpp)
192193
llama_build_and_test(test-batch-alloc.cpp LABEL behavior LINK llama)
193194
llama_build_and_test(test-chat.cpp LABEL behavior WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
194195
target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)

0 commit comments

Comments
 (0)