From e79f2e8051dc6f1532d29b45ffc2f4644ed8da9a Mon Sep 17 00:00:00 2001 From: Michael Johnson Date: Tue, 22 Sep 2026 12:29:09 +0200 Subject: [PATCH 1/2] parallel-decision: expose exact distributions --- tests/CMakeLists.txt | 3 ++ tests/test-parallel-decision.cpp | 38 +++++++++++++++++++++ tools/parallel-decision/README.md | 10 ++++-- tools/parallel-decision/decision-engine.cpp | 10 ++++++ 4 files changed, 58 insertions(+), 3 deletions(-) create mode 100644 tests/test-parallel-decision.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 9b3a4fcc4bbf..dcde2ae9dfbc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -163,6 +163,9 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) + llama_build_and_test(test-parallel-decision.cpp) + target_include_directories(test-parallel-decision PRIVATE ${PROJECT_SOURCE_DIR}/tools/parallel-decision) + target_link_libraries(test-parallel-decision PRIVATE llama-decision) llama_build_and_test(test-json-schema-to-grammar.cpp) if (NOT GGML_BACKEND_DL) diff --git a/tests/test-parallel-decision.cpp b/tests/test-parallel-decision.cpp new file mode 100644 index 000000000000..2f74cb664a23 --- /dev/null +++ b/tests/test-parallel-decision.cpp @@ -0,0 +1,38 @@ +#ifdef NDEBUG +#undef NDEBUG +#endif + +#include "decision-engine.h" + +#include +#include + +static bool near(double actual, double expected) { + return std::fabs(actual - expected) < 1e-6; +} + +int main() { + const common_json schema = common_json::parse(R"({ + "route": { + "type": "enum", + "choices": ["dns", "forgejo", "taskboard"], + "description": "Select a route" + } + })"); + const auto compiled = llama_decision::compile_schema(schema, ""); + + llama_decision::result exact; + exact.fields.push_back({ 1, 0.7f, 2, true, { 0.2f, 0.7f, 0.1f } }); + const common_json assembled = llama_decision::assemble(compiled, exact); + const auto & field = assembled.at("fields").at("route"); + assert(field.at("value") == "forgejo"); + assert(near(field.at("probability").get(), 0.7)); + assert(near(field.at("probabilities").at("dns").get(), 0.2)); + assert(near(field.at("probabilities").at("forgejo").get(), 0.7)); + assert(near(field.at("probabilities").at("taskboard").get(), 0.1)); + + llama_decision::result greedy; + greedy.fields.push_back({ 1, 0.7f, 1, false, {} }); + const common_json partial = llama_decision::assemble(compiled, greedy); + assert(!partial.at("fields").at("route").contains("probabilities")); +} diff --git a/tools/parallel-decision/README.md b/tools/parallel-decision/README.md index 69c0dd1191e7..ac50d8541b7f 100644 --- a/tools/parallel-decision/README.md +++ b/tools/parallel-decision/README.md @@ -84,9 +84,9 @@ curl http://localhost:8096/v1/decision -H "Content-Type: application/json" -d '{ { "decision": {"category": "billing", "urgent": true, "priority": "high"}, "fields": { - "category": {"value": "billing", "probability": 1.0, "scored_nodes": 1, "tree": true}, - "urgent": {"value": true, "probability": 1.0, "scored_nodes": 1, "tree": true}, - "priority": {"value": "high", "probability": 0.74, "scored_nodes": 1, "tree": true} + "category": {"value": "billing", "probability": 1.0, "probabilities": {"billing": 1.0, "technical": 0.0, "cancellation": 0.0, "other": 0.0}, "scored_nodes": 1, "tree": true}, + "urgent": {"value": true, "probability": 1.0, "probabilities": {"true": 1.0, "false": 0.0}, "scored_nodes": 1, "tree": true}, + "priority": {"value": "high", "probability": 0.74, "probabilities": {"low": 0.01, "medium": 0.20, "high": 0.74, "critical": 0.05}, "scored_nodes": 1, "tree": true} }, "usage": {"context_tokens": 21, "scored_rows": 14} } @@ -98,6 +98,10 @@ curl http://localhost:8096/v1/decision -H "Content-Type: application/json" -d '{ (That response is a real one: Gemma 4 12B on an RTX 3060, warm cache.) +In `tree` mode, and in `auto` mode for fields that stay under `tree_max`, each field also returns the complete exact +constrained distribution in `probabilities`. Greedy fields omit it because their unvisited branches do not have exact +probabilities. Keys are the JSON values rendered as strings (`"true"`, `"7"`, `"0.5"`, or the enum value). + ### Schema Compact fields, or a JSON Schema object with `properties`: diff --git a/tools/parallel-decision/decision-engine.cpp b/tools/parallel-decision/decision-engine.cpp index d6ebeb557308..cb00d79b769c 100644 --- a/tools/parallel-decision/decision-engine.cpp +++ b/tools/parallel-decision/decision-engine.cpp @@ -716,6 +716,16 @@ common_json assemble(const compiled_schema & cs, const result & r) { decision[sp.name] = sp.values[idx]; f["value"] = sp.values[idx]; f["probability"] = (double) (fr.probs.size() == sp.values.size() ? fr.probs[idx] : fr.path_score); + if (fr.probs.size() == sp.values.size()) { + common_json probabilities = common_json::object(); + for (size_t k = 0; k < sp.values.size(); ++k) { + const std::string key = sp.values[k].is_string() + ? sp.values[k].get() + : sp.values[k].dump(); + probabilities[key] = (double) fr.probs[k]; + } + f["probabilities"] = std::move(probabilities); + } f["scored_nodes"] = fr.scored_nodes; f["tree"] = fr.tree; fields[sp.name] = f; From f49bb927064b02e06618c4c8937333b9d6c47a2b Mon Sep 17 00:00:00 2001 From: Michael Johnson Date: Tue, 22 Sep 2026 17:16:19 +0200 Subject: [PATCH 2/2] parallel-decision: cache multiple prompt prefixes --- common/arg.cpp | 10 ++++ common/common.h | 1 + tests/test-parallel-decision.cpp | 14 +++++ tools/parallel-decision/README.md | 4 ++ tools/parallel-decision/decision-engine.cpp | 60 ++++++++++++++++++++- tools/parallel-decision/decision-engine.h | 25 ++++++++- tools/server/server-context.cpp | 3 +- 7 files changed, 113 insertions(+), 4 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 5f1211ca40dd..3d280bb401e1 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2557,6 +2557,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.n_seq_decision = value; } ).set_env("LLAMA_ARG_DECISION_SEQS").set_examples({LLAMA_EXAMPLE_SERVER})); + add_opt(common_arg( + {"--decision-prefix-cache"}, "N", + string_format("decision prompt prefixes retained in host memory (default: %d; 0 = active prefix only)", params.n_cache_decision), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("--decision-prefix-cache cannot be negative"); + } + params.n_cache_decision = value; + } + ).set_env("LLAMA_ARG_DECISION_PREFIX_CACHE").set_examples({LLAMA_EXAMPLE_SERVER})); } else { add_opt(common_arg( {"-np", "--parallel"}, "N", diff --git a/common/common.h b/common/common.h index 6c4d2e4e712e..cd85766b9354 100644 --- a/common/common.h +++ b/common/common.h @@ -454,6 +454,7 @@ struct common_params { int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode int32_t n_seq_decision = 0; // sequences reserved for llama-server's /decision endpoint (0 = disabled) + int32_t n_cache_decision = 4; // decision prefix snapshots retained in host memory int32_t n_sequences = 1; // number of sequences to decode int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch) int32_t n_outputs_max_per_seq = 1; // max outputs per sequence diff --git a/tests/test-parallel-decision.cpp b/tests/test-parallel-decision.cpp index 2f74cb664a23..a44cb1b50701 100644 --- a/tests/test-parallel-decision.cpp +++ b/tests/test-parallel-decision.cpp @@ -12,6 +12,20 @@ static bool near(double actual, double expected) { } int main() { + llama_decision::prefix_state_cache cache(2); + cache.put({ 1 }, { 11 }); + cache.put({ 2 }, { 22 }); + assert(cache.size() == 2); + assert(cache.find({ 1 }) && cache.find({ 1 })->at(0) == 11); + cache.put({ 3 }, { 33 }); + assert(cache.size() == 2); + assert(cache.find({ 2 }) == nullptr); // entry 1 was refreshed, so entry 2 was evicted + assert(cache.find({ 1 }) && cache.find({ 3 })); + + llama_decision::prefix_state_cache disabled(0); + disabled.put({ 1 }, { 11 }); + assert(disabled.size() == 0); + const common_json schema = common_json::parse(R"({ "route": { "type": "enum", diff --git a/tools/parallel-decision/README.md b/tools/parallel-decision/README.md index ac50d8541b7f..bfe29d14d371 100644 --- a/tools/parallel-decision/README.md +++ b/tools/parallel-decision/README.md @@ -25,6 +25,10 @@ cmake --build build --config Release -j in flight, the rest are the parallel questions. It also switches the KV cache to unified, which is what lets the branches share the context's cells. +`--decision-prefix-cache N` retains up to N exact prompt-prefix snapshots in host memory (default 4). Alternating +schemas can therefore restore their KV state without re-prefilling or consuming additional decision sequences. Set it +to 0 to keep only the currently active KV prefix. + ```bash ./build/bin/llama-server -m model.gguf -ngl 99 -fa on -c 32768 --decision-seqs 24 --port 8096 ``` diff --git a/tools/parallel-decision/decision-engine.cpp b/tools/parallel-decision/decision-engine.cpp index cb00d79b769c..8fc3ec3f9de0 100644 --- a/tools/parallel-decision/decision-engine.cpp +++ b/tools/parallel-decision/decision-engine.cpp @@ -173,9 +173,43 @@ struct decision_field { // ---------------------------------------------------------------- engine -engine::engine(llama_context * ctx, llama_seq_id seq_base, int n_seqs) +prefix_state_cache::prefix_state_cache(size_t capacity) : capacity(capacity) { +} + +const std::vector * prefix_state_cache::find(const tokens_t & tokens) { + for (auto it = entries.begin(); it != entries.end(); ++it) { + if (it->tokens == tokens) { + entries.splice(entries.begin(), entries, it); + return &entries.front().state; + } + } + return nullptr; +} + +void prefix_state_cache::put(tokens_t tokens, std::vector state) { + if (capacity == 0 || tokens.empty() || state.empty()) { + return; + } + for (auto it = entries.begin(); it != entries.end(); ++it) { + if (it->tokens == tokens) { + it->state = std::move(state); + entries.splice(entries.begin(), entries, it); + return; + } + } + entries.push_front({ std::move(tokens), std::move(state) }); + while (entries.size() > capacity) { + entries.pop_back(); + } +} + +size_t prefix_state_cache::size() const { + return entries.size(); +} + +engine::engine(llama_context * ctx, llama_seq_id seq_base, int n_seqs, size_t prefix_cache_entries) : ctx(ctx), vocab(llama_model_get_vocab(llama_get_model(ctx))), mem(llama_get_memory(ctx)), - seq_snap(seq_base), seq_pool(seq_base + 1), n_pool(n_seqs - 1) { + seq_snap(seq_base), seq_pool(seq_base + 1), n_pool(n_seqs - 1), prefix_cache(prefix_cache_entries) { if (n_seqs < 3) { throw std::invalid_argument("a decision engine needs at least 3 sequences"); } @@ -221,15 +255,37 @@ void engine::decode_parts(const std::vector & parts) { bool engine::prepare_prefix(const tokens_t & shared, bool allow_cache) { if (allow_cache && !shared.empty() && shared == cached && llama_memory_seq_pos_max(mem, seq_snap) == (llama_pos) cached.size() - 1) { + prefix_cache.find(shared); // refresh its LRU position when host snapshots are enabled return true; } for (llama_seq_id s = seq_snap; s < seq_pool + n_pool; ++s) { llama_memory_seq_rm(mem, s, -1, -1); } cached.clear(); + if (allow_cache && !shared.empty()) { + if (const auto * state = prefix_cache.find(shared)) { + const size_t restored = llama_state_seq_set_data_ext( + ctx, state->data(), state->size(), seq_snap, LLAMA_STATE_SEQ_FLAGS_NONE); + if (restored != state->size() || + llama_memory_seq_pos_max(mem, seq_snap) != (llama_pos) shared.size() - 1) { + throw std::runtime_error("failed to restore a cached decision prefix"); + } + cached = shared; + return true; + } + } if (!shared.empty()) { decode_parts({ { &shared, 0, seq_snap } }); cached = shared; + llama_synchronize(ctx); + const size_t state_size = llama_state_seq_get_size_ext(ctx, seq_snap, LLAMA_STATE_SEQ_FLAGS_NONE); + std::vector state(state_size); + const size_t saved = llama_state_seq_get_data_ext( + ctx, state.data(), state.size(), seq_snap, LLAMA_STATE_SEQ_FLAGS_NONE); + if (saved != state.size()) { + throw std::runtime_error("failed to save a decision prefix state"); + } + prefix_cache.put(shared, std::move(state)); } return false; } diff --git a/tools/parallel-decision/decision-engine.h b/tools/parallel-decision/decision-engine.h index c96f5d742269..f033421cf07a 100644 --- a/tools/parallel-decision/decision-engine.h +++ b/tools/parallel-decision/decision-engine.h @@ -13,6 +13,8 @@ #include "llama.h" #include "json.h" +#include +#include #include #include #include @@ -23,6 +25,26 @@ namespace llama_decision { using tokens_t = std::vector; +// Small exact-key LRU used by the engine to retain serialized prefix states in host memory. +// Keeping snapshots off the KV sequence pool preserves every reserved sequence for scoring. +class prefix_state_cache { + public: + explicit prefix_state_cache(size_t capacity); + + const std::vector * find(const tokens_t & tokens); + void put(tokens_t tokens, std::vector state); + size_t size() const; + + private: + struct entry { + tokens_t tokens; + std::vector state; + }; + + size_t capacity; + std::list entries; // most recently used first +}; + // One field as the scorer sees it: the text before its value and the allowed value texts. struct field_input { std::string suffix; // e.g. ' "fire": ' @@ -72,7 +94,7 @@ struct batch_result { // flight, then branches. The context needs a unified KV cache so branches share the trunk's cells. class engine { public: - engine(llama_context * ctx, llama_seq_id seq_base, int n_seqs); + engine(llama_context * ctx, llama_seq_id seq_base, int n_seqs, size_t prefix_cache_entries = 4); result decide(const std::string & shared_text, const std::string & context_text, const std::vector & fields, const options & opt); @@ -101,6 +123,7 @@ class engine { llama_seq_id seq_snap, seq_pool; int n_pool; tokens_t cached; + prefix_state_cache prefix_cache; tokens_t tokenize(const std::string & text, bool add_special) const; void decode_parts(const std::vector & parts); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 41c6da3e1bbc..639563afc6b0 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2399,7 +2399,8 @@ struct server_context_impl { } if (!decision_engine) { decision_engine = std::make_unique(ctx_tgt, (llama_seq_id) params_base.n_parallel, - params_base.n_seq_decision); + params_base.n_seq_decision, + (size_t) params_base.n_cache_decision); } const auto cs = llama_decision::compile_schema(body.at("schema"), body.value("instructions", std::string())); std::string shared;