Suffix decode

common/CMakeLists.txt

@@ -100,6 +100,8 @@ add_library(${TARGET}
     sampling.h
     speculative.cpp
     speculative.h
+    suffix-tree.cpp
+    suffix-tree.h
     subproc.cpp
     subproc.h
     trie.cpp

common/arg.cpp

@@ -4208,6 +4208,51 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         }
     ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
 
+    add_opt(common_arg(
+        {"--spec-ngram-suffix-max-depth"}, "N",
+        string_format("suffix-tree depth = max context-match length for ngram-suffix speculative decoding (default: %d)", params.speculative.ngram_suffix.max_depth),
+        [](common_params & params, int value) {
+            if (value < 1) {
+                throw std::invalid_argument("ngram-suffix max depth must be at least 1");
+            }
+            params.speculative.ngram_suffix.max_depth = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-ngram-suffix-n-max"}, "N",
+        string_format("maximum number of draft tokens for ngram-suffix speculative decoding (default: %d)", params.speculative.ngram_suffix.n_max),
+        [](common_params & params, int value) {
+            if (value < 1) {
+                throw std::invalid_argument("ngram-suffix n-max must be at least 1");
+            }
+            params.speculative.ngram_suffix.n_max = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-ngram-suffix-n-min"}, "N",
+        string_format("discard drafts shorter than this for ngram-suffix speculative decoding (default: %d)", params.speculative.ngram_suffix.n_min),
+        [](common_params & params, int value) {
+            if (value < 1) {
+                throw std::invalid_argument("ngram-suffix n-min must be at least 1");
+            }
+            params.speculative.ngram_suffix.n_min = value;
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-ngram-suffix-max-factor"}, "F",
+        string_format("draft up to match_len * F tokens for ngram-suffix speculative decoding (default: %.1f)", params.speculative.ngram_suffix.max_factor),
+        [](common_params & params, const std::string & value) {
+            params.speculative.ngram_suffix.max_factor = std::stof(value);
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--spec-ngram-suffix-min-prob"}, "F",
+        string_format("stop drafting below this frequency probability for ngram-suffix speculative decoding (default: %.2f)", params.speculative.ngram_suffix.min_prob),
+        [](common_params & params, const std::string & value) {
+            params.speculative.ngram_suffix.min_prob = std::stof(value);
+        }
+    ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+
     //
     // removed params
     //

common/common.h

@@ -179,6 +179,7 @@ enum common_speculative_type {
     COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values
     COMMON_SPECULATIVE_TYPE_NGRAM_MOD,
     COMMON_SPECULATIVE_TYPE_NGRAM_CACHE,   // self-speculative decoding with 3-level n-gram cache
+    COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX,  // self-speculative decoding with a suffix tree (suffix decoding)
     COMMON_SPECULATIVE_TYPE_COUNT          // number of types, unknown type
 };
 
@@ -366,6 +367,14 @@ struct common_params_speculative_ngram_cache {
     std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding
 };
 
+struct common_params_speculative_ngram_suffix {
+    int32_t max_depth = 24;    // suffix-tree depth = max context-match length
+    int32_t n_max     = 24;    // maximum number of drafted tokens
+    int32_t n_min     = 3;     // discard drafts shorter than this
+    float   max_factor = 1.0f; // draft up to match_len * max_factor tokens
+    float   min_prob   = 0.1f; // stop drafting below this frequency probability
+};
+
 struct common_params_speculative {
     std::vector<enum common_speculative_type> types = { COMMON_SPECULATIVE_TYPE_NONE };
 
@@ -378,6 +387,7 @@ struct common_params_speculative {
     common_params_speculative_ngram_map ngram_map_k4v;
 
     common_params_speculative_ngram_cache ngram_cache;
+    common_params_speculative_ngram_suffix ngram_suffix;
 
     bool has_dft() const {
         return !draft.mparams.empty();

common/speculative.cpp

@@ -8,6 +8,7 @@
 #include "ngram-map.h"
 #include "ngram-mod.h"
 #include "sampling.h"
+#include "suffix-tree.h"
 
 #include "../src/llama-ext.h" // staging API: llama_set_embeddings_nextn / llama_get_embeddings_nextn_ith (used by MTP)
 
@@ -39,7 +40,8 @@ const std::map<std::string, common_speculative_type> common_speculative_type_fro
     {"ngram-map-k",   COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
     {"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
     {"ngram-mod",     COMMON_SPECULATIVE_TYPE_NGRAM_MOD},
-    {"ngram-cache",   COMMON_SPECULATIVE_TYPE_NGRAM_CACHE}
+    {"ngram-cache",   COMMON_SPECULATIVE_TYPE_NGRAM_CACHE},
+    {"ngram-suffix",  COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX}
 };
 
 static std::string common_speculative_get_devices_str(const std::vector<ggml_backend_dev_t> & devices) {
@@ -1742,6 +1744,102 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
     }
 };
 
+// suffix decoding: a model-free drafter backed by a per-sequence suffix tree
+// built over the prompt and previously generated tokens.
+struct common_speculative_impl_ngram_suffix : public common_speculative_impl {
+    common_params_speculative_ngram_suffix params;
+
+    // per-sequence suffix tree plus how many committed tokens are already in it.
+    struct seq_state {
+        common_suffix_tree tree;
+        size_t n_inserted = 0;
+
+        explicit seq_state(int32_t max_depth) : tree(max_depth) {}
+    };
+    std::vector<seq_state> states;
+
+    common_speculative_impl_ngram_suffix(const common_params_speculative & params, uint32_t n_seq)
+        : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX, n_seq)
+        , params(params.ngram_suffix)
+    {
+        SPC_TRC("%s", "adding speculative implementation 'ngram-suffix'\n");
+        SPC_TRC("- max_depth=%d, n_max=%d, max_factor=%f, min_prob=%f\n",
+                this->params.max_depth, this->params.n_max, this->params.max_factor, this->params.min_prob);
+
+        states.reserve(n_seq);
+        for (uint32_t i = 0; i < n_seq; ++i) {
+            states.emplace_back(this->params.max_depth);
+        }
+    }
+
+    void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
+        GGML_ASSERT(seq_id < (llama_seq_id) n_seq);
+
+        auto & st = states[seq_id];
+        st.tree.reset();
+        st.tree.extend(prompt);
+        st.n_inserted = prompt.size();
+    }
+
+    bool process(const llama_batch & /*batch*/) override {
+        // TODO: implement
+        return true;
+    }
+
+    void draft(common_speculative_draft_params_vec & dparams) override {
+        assert(dparams.size() == n_seq);
+
+        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+            auto & dp = dparams[seq_id];
+            if (!dp.drafting) {
+                continue;
+            }
+
+            auto & st = states[seq_id];
+            const llama_tokens & prompt = *dp.prompt;
+
+            // sync the tree with the committed prompt (accepted tokens only, so
+            // it only ever grows - rejected drafts never reach it).
+            if (prompt.size() < st.n_inserted) {
+                // history was truncated (e.g. context shift) -> rebuild.
+                st.tree.reset();
+                st.tree.extend(prompt);
+            } else {
+                for (size_t i = st.n_inserted; i < prompt.size(); ++i) {
+                    st.tree.append(prompt[i]);
+                }
+            }
+            st.n_inserted = prompt.size();
+
+            // context = tail of the committed prompt + the just-sampled token
+            // (which is not yet part of the tree).
+            llama_tokens context;
+            const int32_t take = std::min<int32_t>((int32_t) prompt.size(), params.max_depth - 1);
+            context.reserve(take + 1);
+            context.insert(context.end(), prompt.end() - take, prompt.end());
+            context.push_back(dp.id_last);
+
+            common_suffix_draft draft = st.tree.speculate(
+                    context, params.n_max, params.max_factor, params.min_prob);
+
+            // a short draft is unlikely to pay for its verification batch.
+            if ((int32_t) draft.tokens.size() < params.n_min) {
+                continue;
+            }
+
+            *dp.result = std::move(draft.tokens);
+        }
+    }
+
+    void accept(llama_seq_id /*seq_id*/, uint16_t /*n_accepted*/, bool /*is_other*/) override {
+        // noop
+    }
+
+    bool need_embd() const override {
+        return false;
+    }
+};
+
 struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
     // n_seq configs
     std::vector<common_ngram_map> config;
@@ -2196,6 +2294,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) {
         case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
         case COMMON_SPECULATIVE_TYPE_NGRAM_MOD:     return "ngram-mod";
         case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE:   return "ngram-cache";
+        case COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX:  return "ngram-suffix";
         default:                                    return "unknown";
     }
 }
@@ -2262,6 +2361,9 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
             case COMMON_SPECULATIVE_TYPE_NGRAM_CACHE:
                 n_max = std::max(n_max, (int32_t) 8);
                 break;
+            case COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX:
+                n_max = std::max(n_max, spec->ngram_suffix.n_max);
+                break;
             case COMMON_SPECULATIVE_TYPE_NONE:
             case COMMON_SPECULATIVE_TYPE_COUNT:
                 break;
@@ -2392,7 +2494,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         };
 
         // when adding a new type - update here the logic above
-        static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11);
+        static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 12);
 
         // this list here defines the priority of the speculators
         // the one with highest priority are listed first
@@ -2401,6 +2503,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
         add_config_if_enabled(COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V);
         add_config_if_enabled(COMMON_SPECULATIVE_TYPE_NGRAM_MOD);
         add_config_if_enabled(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE);
+        add_config_if_enabled(COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX);
 
         add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE);
         add_config_if_enabled(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, params.draft.ctx_dft != nullptr);
@@ -2479,6 +2582,10 @@ common_speculative * common_speculative_init(common_params_speculative & params,
                 impls.push_back(std::make_unique<common_speculative_impl_ngram_cache>(state));
                 break;
             }
+            case COMMON_SPECULATIVE_TYPE_NGRAM_SUFFIX: {
+                impls.push_back(std::make_unique<common_speculative_impl_ngram_suffix>(config.params, n_seq));
+                break;
+            }
             default:
                 break;
         }