1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
#pragma once
#include "llama.h"
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
// A trie-based suffix tree for suffix-decoding speculative decoding.
//
// Stores all suffixes (up to max_depth) of the token history.
// Used to find matching patterns in context and generate draft tokens
// by following the most frequent continuation path.
//
// Reference: "Suffix Decoding" (Saxena et al., 2024) — arXiv:2411.04975
struct common_suffix_node {
int64_t count = 0;
std::unordered_map<llama_token, std::unique_ptr<common_suffix_node>> children;
};
class common_suffix_tree {
public:
explicit common_suffix_tree(int max_depth = 64);
~common_suffix_tree();
// Append tokens to the history and insert new suffixes into the trie.
// Incremental: only processes suffixes that haven't been inserted yet.
void extend(const llama_token * tokens, int n_tokens);
void clear();
// Generate draft tokens by matching the context in the trie.
// Tries multiple context lengths and returns the draft with the best score.
std::vector<llama_token> speculate(
const llama_token * context, int n_context,
int max_spec_tokens,
float min_token_prob = 0.1f,
int min_match_count = 1,
int min_match_len = 5) const;
// Load an offline corpus to pre-warm the tree before any request.
// Supported formats (.json or .bin)
bool load_corpus(
const std::string & path,
std::function<std::vector<llama_token>(const std::string &)> tokenize_fn = {});
int max_depth() const { return _max_depth; }
int token_count() const { return (int)_tokens.size(); }
private:
int _max_depth;
std::unique_ptr<common_suffix_node> _root;
std::vector<llama_token> _tokens;
int _n_inserted = 0;
void _insert_suffix(int start_pos);
void _extend_suffix(int start_pos, int old_len, int new_len);
};