parallel-decision: expose exact tree distributions

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",

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

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)

tests/test-parallel-decision.cpp

@@ -0,0 +1,52 @@
+#ifdef NDEBUG
+#undef NDEBUG
+#endif
+
+#include "decision-engine.h"
+
+#include <cassert>
+#include <cmath>
+
+static bool near(double actual, double expected) {
+    return std::fabs(actual - expected) < 1e-6;
+}
+
+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",
+            "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<double>(), 0.7));
+    assert(near(field.at("probabilities").at("dns").get<double>(), 0.2));
+    assert(near(field.at("probabilities").at("forgejo").get<double>(), 0.7));
+    assert(near(field.at("probabilities").at("taskboard").get<double>(), 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"));
+}

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
 ```
@@ -84,9 +88,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 +102,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`:

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<uint8_t> * 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<uint8_t> 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<prompt_part> & 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<uint8_t> 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;
 }
@@ -716,6 +772,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<std::string>()
+                    : 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;

tools/parallel-decision/decision-engine.h

@@ -13,6 +13,8 @@
 #include "llama.h"
 #include "json.h"
 
+#include <cstddef>
+#include <list>
 #include <string>
 #include <utility>
 #include <vector>
@@ -23,6 +25,26 @@ namespace llama_decision {
 
 using tokens_t = std::vector<llama_token>;
 
+// 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<uint8_t> * find(const tokens_t & tokens);
+    void put(tokens_t tokens, std::vector<uint8_t> state);
+    size_t size() const;
+
+  private:
+    struct entry {
+        tokens_t            tokens;
+        std::vector<uint8_t> state;
+    };
+
+    size_t           capacity;
+    std::list<entry> 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<field_input> & 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<prompt_part> & parts);

tools/server/server-context.cpp

@@ -2399,7 +2399,8 @@ struct server_context_impl {
         }
         if (!decision_engine) {
             decision_engine = std::make_unique<llama_decision::engine>(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;