ggml-vulkan: add optional device keepalive (#29267)

* ggml-vulkan: add optional device keepalive

* simplify implementation, move ownership to backend, change env var to seconds

---------

Co-authored-by: Acmmi <acmmi@Acmmis-MacBook-Air.local>
Co-authored-by: Ruben Ortlam <rortlam@redhat.com>

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -1279,6 +1279,10 @@ struct ggml_backend_vk_context {
 
     vk_device device;
 
+    // Workaround for drivers moving VRAM contents off the GPU when idle
+    std::promise<void> keepalive_stop;
+    std::thread keepalive_thread;
+
     size_t semaphore_idx, binary_semaphore_idx, event_idx;
     ggml_vk_garbage_collector gc;
     size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset;

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -4237,6 +4237,8 @@ static bool ggml_vk_has_rebar(const vk::PhysicalDevice & physical_device) {
     return false;
 }
 
+static void ggml_vk_device_start_keepalive(vk_device& device);
+
 vk_device ggml_vk_get_device(size_t idx) {
     VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")");
 
@@ -13353,6 +13355,10 @@ void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) {
 
 void ggml_vk_cleanup(ggml_backend_vk_context * ctx) {
     VK_LOG_DEBUG("ggml_vk_cleanup(" << ctx->name << ")");
+    if (ctx->keepalive_thread.joinable()) {
+        ctx->keepalive_stop.set_value();
+        ctx->keepalive_thread.join();
+    }
     // discard any unsubmitted command buffers
     ctx->compute_ctx.reset();
     // wait for any pending command buffers to finish
@@ -15719,11 +15725,34 @@ static ggml_guid_t ggml_backend_vk_guid() {
     return &guid;
 }
 
+static void ggml_vk_start_keepalive(ggml_backend_vk_context * ctx) {
+    const char * env = getenv("GGML_VK_KEEPALIVE");
+    const int s = env ? atoi(env) : 0;
+    if (s <= 0) {
+        return;
+    }
+
+    GGML_LOG_INFO("ggml_vulkan: GGML_VK_KEEPALIVE=%ds on %s\n", s, ctx->device->name.c_str());
+
+    vk_buffer buf = ggml_vk_create_buffer_check(ctx->device, 4, vk::MemoryPropertyFlagBits::eDeviceLocal);
+    ctx->keepalive_thread = std::thread([buf, stop = ctx->keepalive_stop.get_future(), s]() mutable {
+        while (stop.wait_for(std::chrono::seconds(s)) == std::future_status::timeout) {
+            try {
+                ggml_vk_buffer_memset(buf, 0, 0, 4);
+            } catch (const std::exception & e) {
+                GGML_LOG_WARN("ggml_vulkan: keepalive stopped: %s\n", e.what());
+                return;
+            }
+        }
+    });
+}
+
 ggml_backend_t ggml_backend_vk_init(size_t dev_num) {
     VK_LOG_DEBUG("ggml_backend_vk_init(" << dev_num << ")");
 
     ggml_backend_vk_context * ctx = new ggml_backend_vk_context;
     ggml_vk_init(ctx, dev_num);
+    ggml_vk_start_keepalive(ctx);
 
     ggml_backend_t vk_backend = new ggml_backend {
         /* .guid    = */ ggml_backend_vk_guid(),

src/llama-grammar.cpp

@@ -948,6 +948,160 @@ static void llama_grammar_advance_stack(
     }
 }
 
+// Left-factor grammar rules to reduce redundant parallel stacks.
+//
+// When a rule has many alternates (e.g. tool-choice ::= tc0 | tc1 | ... | tc119)
+// and the referenced rules share a common element-content prefix, extract that
+// prefix so the parser processes it once instead of N times.
+//
+// Handles two cases:
+// 1. Inline common prefix: rule = "ab1" | "ab2" -> rule = "ab" rule', rule' = "1" | "2"
+// 2. Cross-rule prefix: rule = A | B | C where A,B,C start with same elements
+//    -> rule = common_prefix rule', rule' = A_suffix | B_suffix | C_suffix
+static void llama_grammar_left_factor(llama_grammar_rules & rules) {
+    auto elem_eq = [](const llama_grammar_element & a, const llama_grammar_element & b) {
+        return a.type == b.type && a.value == b.value;
+    };
+
+    // Keep the prefix boundary outside character classes.
+    auto trim_prefix = [](const llama_grammar_rule & rule, size_t start, size_t length) {
+        while (length > 0) {
+            const auto type = rule[start + length].type;
+            if (type != LLAMA_GRETYPE_CHAR_ALT && type != LLAMA_GRETYPE_CHAR_RNG_UPPER) {
+                break;
+            }
+            length--;
+        }
+        return length;
+    };
+
+    struct alt_info {
+        size_t rule_idx;
+        size_t start;
+        size_t end;
+    };
+
+    const size_t n_rules_original = rules.size();
+    // Original rules factor at most once; generated rules do not resolve references.
+    // Each new suffix removes a non-empty prefix.
+    bool changed = true;
+    while (changed) {
+        changed = false;
+        const size_t n_rules = rules.size();
+
+        for (size_t ri = 0; ri < n_rules; ri++) {
+            const auto & rule = rules[ri];
+
+            // Split rule into alternates
+            std::vector<std::pair<size_t, size_t>> alt_ranges;
+            size_t alt_start = 0;
+            for (size_t j = 0; j < rule.size(); j++) {
+                if (rule[j].type == LLAMA_GRETYPE_ALT || rule[j].type == LLAMA_GRETYPE_END) {
+                    alt_ranges.push_back({alt_start, j});
+                    alt_start = j + 1;
+                }
+            }
+
+            if (alt_ranges.size() < 4) {
+                continue;
+            }
+
+            // Resolve each alternate: if it's a single RULE_REF, follow it to get the
+            // actual element sequence from the referenced rule.
+            std::vector<alt_info> resolved;
+            resolved.reserve(alt_ranges.size());
+            bool all_single_refs = true;
+
+            for (const auto & [astart, aend] : alt_ranges) {
+                size_t len = aend - astart;
+                // Resolving references in generated suffixes can expand recursive rules forever.
+                if (ri < n_rules_original && len == 1 && rule[astart].type == LLAMA_GRETYPE_RULE_REF) {
+                    size_t ref_id = static_cast<size_t>(rule[astart].value);
+                    if (ref_id < rules.size()) {
+                        size_t ref_end = 0;
+                        for (size_t k = 0; k < rules[ref_id].size(); k++) {
+                            if (rules[ref_id][k].type == LLAMA_GRETYPE_ALT ||
+                                rules[ref_id][k].type == LLAMA_GRETYPE_END) {
+                                ref_end = k;
+                                break;
+                            }
+                        }
+                        bool single_alt = (rules[ref_id][ref_end].type == LLAMA_GRETYPE_END);
+                        if (single_alt) {
+                            resolved.push_back({ref_id, 0, ref_end});
+                            continue;
+                        }
+                    }
+                }
+                all_single_refs = false;
+                resolved.push_back({ri, astart, aend});
+            }
+
+            if (!all_single_refs) {
+                resolved.clear();
+                for (const auto & [astart, aend] : alt_ranges) {
+                    resolved.push_back({ri, astart, aend});
+                }
+            }
+
+            size_t prefix_len = 0;
+            bool prefix_match = true;
+            while (prefix_match) {
+                const auto & first = resolved[0];
+                if (first.start + prefix_len >= first.end) {
+                    break;
+                }
+                const auto & ref_elem = rules[first.rule_idx][first.start + prefix_len];
+
+                for (size_t ai = 1; ai < resolved.size(); ai++) {
+                    const auto & alt = resolved[ai];
+                    if (alt.start + prefix_len >= alt.end ||
+                        !elem_eq(rules[alt.rule_idx][alt.start + prefix_len], ref_elem)) {
+                        prefix_match = false;
+                        break;
+                    }
+                }
+                if (prefix_match) {
+                    prefix_len++;
+                }
+            }
+
+            for (const auto & alt : resolved) {
+                prefix_len = trim_prefix(rules[alt.rule_idx], alt.start, prefix_len);
+            }
+            if (prefix_len == 0) {
+                continue;
+            }
+
+            uint32_t suffix_choice_id = static_cast<uint32_t>(rules.size());
+            llama_grammar_rule suffix_choice;
+
+            for (size_t ai = 0; ai < resolved.size(); ai++) {
+                if (ai > 0) {
+                    suffix_choice.push_back({LLAMA_GRETYPE_ALT, 0});
+                }
+                const auto & alt = resolved[ai];
+                for (size_t j = alt.start + prefix_len; j < alt.end; j++) {
+                    suffix_choice.push_back(rules[alt.rule_idx][j]);
+                }
+            }
+            suffix_choice.push_back({LLAMA_GRETYPE_END, 0});
+
+            llama_grammar_rule new_rule;
+            const auto & first = resolved[0];
+            for (size_t j = 0; j < prefix_len; j++) {
+                new_rule.push_back(rules[first.rule_idx][first.start + j]);
+            }
+            new_rule.push_back({LLAMA_GRETYPE_RULE_REF, suffix_choice_id});
+            new_rule.push_back({LLAMA_GRETYPE_END, 0});
+
+            rules.push_back(std::move(suffix_choice));
+            rules[ri] = std::move(new_rule);
+            changed = true;
+        }
+    }
+}
+
 static bool llama_grammar_detect_left_recursion(
         const llama_grammar_rules & rules,
         size_t rule_index,
@@ -1006,7 +1160,7 @@ static bool llama_grammar_detect_left_recursion(
 }
 
 const llama_grammar_rules & llama_grammar_get_rules(const struct llama_grammar * grammar) {
-    return grammar->rules;
+    return *grammar->rules;
 }
 
 llama_grammar_stacks & llama_grammar_get_stacks(struct llama_grammar * grammar) {
@@ -1035,7 +1189,7 @@ static void llama_grammar_accept_chr(
         if (!llama_grammar_is_end_of_sequence(match.second)) {
             new_stack.push_back(match.second);
         }
-        llama_grammar_advance_stack(grammar.rules, new_stack, new_stacks);
+        llama_grammar_advance_stack(*grammar.rules, new_stack, new_stacks);
     }
 }
 
@@ -1221,12 +1375,16 @@ struct llama_grammar * llama_grammar_init_impl(
         }
     }
 
-    // Check for left recursion
-    std::vector<bool> rules_visited(n_rules);
-    std::vector<bool> rules_in_progress(n_rules);
-    std::vector<bool> rules_may_be_empty(n_rules);
+    // Left-factor rules to reduce redundant parallel stacks
+    llama_grammar_left_factor(vec_rules);
+
+    // Check for left recursion (use post-factoring rule count)
+    const size_t n_rules_total = vec_rules.size();
+    std::vector<bool> rules_visited(n_rules_total);
+    std::vector<bool> rules_in_progress(n_rules_total);
+    std::vector<bool> rules_may_be_empty(n_rules_total);
     try {
-        for (size_t i = 0; i < n_rules; i++) {
+        for (size_t i = 0; i < n_rules_total; i++) {
             if (rules_visited[i]) {
                 continue;
             }
@@ -1267,7 +1425,7 @@ struct llama_grammar * llama_grammar_init_impl(
     // then the pointers would be invalidated when the local vec_rules goes out of scope.
     return new llama_grammar {
         vocab,
-        std::move(vec_rules),
+        std::make_shared<const llama_grammar_rules>(std::move(vec_rules)),
         std::move(stacks),
         /* .partial_utf8 = */             {},
         /* .lazy = */                     false,
@@ -1319,12 +1477,16 @@ struct llama_grammar * llama_grammar_init_impl(
         vec_rules[i].push_back({LLAMA_GRETYPE_END, 0});
     }
 
-    // Check for left recursion
-    std::vector<bool> rules_visited(n_rules);
-    std::vector<bool> rules_in_progress(n_rules);
-    std::vector<bool> rules_may_be_empty(n_rules);
+    // Left-factor rules to reduce redundant parallel stacks
+    llama_grammar_left_factor(vec_rules);
+
+    // Check for left recursion (use post-factoring rule count)
+    const size_t n_rules_total = vec_rules.size();
+    std::vector<bool> rules_visited(n_rules_total);
+    std::vector<bool> rules_in_progress(n_rules_total);
+    std::vector<bool> rules_may_be_empty(n_rules_total);
     try {
-        for (size_t i = 0; i < n_rules; i++) {
+        for (size_t i = 0; i < n_rules_total; i++) {
             if (rules_visited[i]) {
                 continue;
             }
@@ -1378,7 +1540,7 @@ struct llama_grammar * llama_grammar_init_impl(
     // then the pointers would be invalidated when the local vec_rules goes out of scope.
     return new llama_grammar {
         vocab,
-        std::move(vec_rules),
+        std::make_shared<const llama_grammar_rules>(std::move(vec_rules)),
         std::move(stacks),
         /* .partial_utf8 = */             {},
         /* .lazy = */                     lazy,
@@ -1399,33 +1561,7 @@ void llama_grammar_free_impl(struct llama_grammar * grammar) {
 }
 
 struct llama_grammar * llama_grammar_clone_impl(const struct llama_grammar & grammar) {
-    auto * result = new llama_grammar {
-        grammar.vocab,
-        grammar.rules,
-        grammar.stacks,
-        grammar.partial_utf8,
-        grammar.lazy,
-        grammar.awaiting_trigger,
-        grammar.trigger_buffer,
-        grammar.trigger_buffer_positions,
-        grammar.trigger_tokens,
-        grammar.trigger_patterns,
-    };
-
-    // redirect elements in stacks to point to new rules
-    for (size_t is = 0; is < result->stacks.size(); is++) {
-        for (size_t ie = 0; ie < result->stacks[is].size(); ie++) {
-            for (size_t ir0 = 0; ir0 < grammar.rules.size(); ir0++) {
-                for (size_t ir1 = 0; ir1 < grammar.rules[ir0].size(); ir1++) {
-                    if (grammar.stacks[is][ie] == &grammar.rules[ir0][ir1]) {
-                        result->stacks[is][ie] =  &result->rules[ir0][ir1];
-                    }
-                }
-            }
-        }
-    }
-
-    return result;
+    return new llama_grammar(grammar);
 }
 
 void llama_grammar_apply_impl(const struct llama_grammar & grammar, llama_token_data_array * cur_p) {
@@ -1465,7 +1601,7 @@ void llama_grammar_apply_impl(const struct llama_grammar & grammar, llama_token_
         }
     }
 
-    const auto rejects = llama_grammar_reject_candidates(grammar.rules, grammar.stacks, candidates_grammar);
+    const auto rejects = llama_grammar_reject_candidates(*grammar.rules, grammar.stacks, candidates_grammar);
     for (const auto & reject : rejects) {
         cur_p->data[reject.index].logit = -INFINITY;
     }
@@ -1566,7 +1702,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);
             }
         } else {
             llama_grammar_stacks current_stacks = {stack};

src/llama-grammar.h

@@ -3,6 +3,7 @@
 #include "llama.h"
 
 #include <map>
+#include <memory>
 #include <regex>
 #include <string>
 #include <vector>
@@ -132,8 +133,8 @@ struct llama_grammar {
     // note: allow null vocab for testing (not great)
     const llama_vocab * vocab;
 
-    const llama_grammar_rules  rules;  // TODO: shared ptr
-          llama_grammar_stacks stacks;
+    const std::shared_ptr<const llama_grammar_rules> rules;
+    llama_grammar_stacks stacks;
 
     // buffer for partially generated UTF-8 sequence from accepted tokens
     llama_partial_utf8 partial_utf8;

tests/test-grammar-integration.cpp

@@ -193,6 +193,19 @@ static void test_schema(const std::string & test_desc, const std::string & schem
 }
 
 static void test_simple_grammar() {
+    test_grammar(
+        "left factor leading empty suffix",
+        "root ::= \"ab\" | \"abc\" | \"abd\" | \"abe\"",
+        {"ab", "abc", "abd", "abe"}, {"", "a", "abb", "abcc"});
+    test_grammar(
+        "left factor middle empty suffix",
+        "root ::= \"abc\" | \"ab\" | \"abd\" | \"abe\"",
+        {"ab", "abc", "abd", "abe"}, {"", "a", "abb", "abcc"});
+    test_grammar(
+        "left factor recursive references",
+        "root ::= \"ok\" | choice\nchoice ::= a | b | c | d\na ::= \"x\" a\nb ::= \"x\" b\nc ::= \"x\" c\nd ::= \"x\" d",
+        {"ok"}, {"", "x", "xx", "okx"});
+
     test_schema(
         "min 0",
         R"""({
@@ -623,6 +636,19 @@ static void test_special_chars() {
             "🔵🟠abc🟠🔵"
         }
     );
+
+    test_grammar(
+        "skip left factoring inside character classes",
+        "root ::= [ab] \"x\" | [ac] \"x\" | [ad] \"x\" | [ae] \"x\"",
+        {"ax", "bx", "cx", "dx", "ex"}, {"", "a", "fx", "axx"});
+    test_grammar(
+        "skip left factoring inside character ranges",
+        "root ::= [a-b] \"x\" | [a-c] \"x\" | [a-d] \"x\" | [a-e] \"x\"",
+        {"ax", "bx", "cx", "dx", "ex"}, {"", "a", "fx", "axx"});
+    test_grammar(
+        "left factor referenced character alternatives",
+        "root ::= a | b | c | d\na ::= \"p\" [a] \"x\"\nb ::= \"p\" [ab] \"x\"\nc ::= \"p\" [ac] \"x\"\nd ::= \"p\" [ad] \"x\"",
+        {"pax", "pbx", "pcx", "pdx"}, {"", "px", "pex", "paxx"});
 }
 
 static void test_quantifiers() {