diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 730a22745..961d45505 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -298,7 +298,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--speculative-ngram-max-match-window-size` | The maximum window size for pattern matching in ngram speculative decoding. | `12` | Type: int | | `--speculative-ngram-min-bfs-breadth` | The minimum breadth for BFS (Breadth-First Search) in ngram speculative decoding. | `1` | Type: int | | `--speculative-ngram-max-bfs-breadth` | The maximum breadth for BFS (Breadth-First Search) in ngram speculative decoding. | `10` | Type: int | -| `--speculative-ngram-match-type` | The match type for cache tree. | `BFS` | `BFS`, `PROB` | +| `--speculative-ngram-match-type` | Ngram tree-building mode. `BFS` selects recency-based expansion and `PROB` selects frequency-based expansion. This setting is forwarded to the ngram cache implementation. | `BFS` | `BFS`, `PROB` | | `--speculative-ngram-branch-length` | The branch length for ngram speculative decoding. | `18` | Type: int | | `--speculative-ngram-capacity` | The cache capacity for ngram speculative decoding. | `10000000` | Type: int | diff --git a/docs/advanced_features/speculative_decoding.md b/docs/advanced_features/speculative_decoding.md index 121ae8d58..0ea688d00 100644 --- a/docs/advanced_features/speculative_decoding.md +++ b/docs/advanced_features/speculative_decoding.md @@ -392,7 +392,7 @@ Enable it with: | `--speculative-ngram-max-match-window-size` | Maximum matching window size. | `12` | | `--speculative-ngram-min-bfs-breadth` | Minimum BFS breadth. | `1` | | `--speculative-ngram-max-bfs-breadth` | Maximum BFS breadth. | `10` | -| `--speculative-ngram-match-type` | Match type: `"BFS"` or `"PROB"`. | `"BFS"` | +| `--speculative-ngram-match-type` | Ngram tree-building mode: `"BFS"` for recency-based expansion or `"PROB"` for frequency-based expansion. | `"BFS"` | | `--speculative-ngram-branch-length` | How many recent tokens to insert into the cache. | `18` | | `--speculative-ngram-capacity` | Cache capacity (number of entries). | `10,000,000` | @@ -468,7 +468,7 @@ Below is a comprehensive list of all speculative decoding parameters available i | `--speculative-ngram-max-match-window-size` | `int` | `12` | Maximum ngram matching window | | `--speculative-ngram-min-bfs-breadth` | `int` | `1` | Minimum BFS breadth | | `--speculative-ngram-max-bfs-breadth` | `int` | `10` | Maximum BFS breadth | -| `--speculative-ngram-match-type` | `str` | `"BFS"` | Match type: `"BFS"` or `"PROB"` | +| `--speculative-ngram-match-type` | `str` | `"BFS"` | Ngram tree-building mode: `"BFS"` for recency-based expansion or `"PROB"` for frequency-based expansion | | `--speculative-ngram-branch-length` | `int` | `18` | Recent tokens to insert into cache | | `--speculative-ngram-capacity` | `int` | `10,000,000` | Cache capacity | diff --git a/docs/platforms/ascend_npu_support_features.md b/docs/platforms/ascend_npu_support_features.md index 3297b6720..4a6e9eeb5 100644 --- a/docs/platforms/ascend_npu_support_features.md +++ b/docs/platforms/ascend_npu_support_features.md @@ -228,7 +228,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen | `--speculative-ngram-`
`max-match-window-size` | `12` | Type: int | Experimental | | `--speculative-ngram-`
`min-bfs-breadth` | `1` | Type: int | Experimental | | `--speculative-ngram-`
`max-bfs-breadth` | `10` | Type: int | Experimental | -| `--speculative-ngram-`
`match-type` | `BFS` | `BFS`,
`PROB` | Experimental | +| `--speculative-ngram-`
`match-type` | `BFS` | `BFS`,
`PROB` | Experimental. `BFS` uses recency-based expansion; `PROB` uses frequency-based expansion. | | `--speculative-ngram-`
`branch-length` | `18` | Type: int | Experimental | | `--speculative-ngram-`
`capacity` | `10000000` | Type: int | Experimental | diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp index e7f0297e2..a5ca8c840 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp +++ b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp @@ -1,74 +1,16 @@ #include "ngram.h" -#include #include -#include #include -#include -#include -#include #include +#include #include -#include -#include -#include + +#include "trie.h" namespace ngram { -struct Node { - std::unordered_map next; -}; - -Ngram::Result fillResult(int last_token, int draft_token_num, std::vector& tree, int root) { - Ngram::Result info; - std::vector prevs; - info.token.reserve(draft_token_num); - prevs.reserve(draft_token_num); - std::queue> queue; - info.token.emplace_back(last_token); - prevs.emplace_back(-1); - - for (auto [token, next] : tree[root].next) { - queue.emplace(token, next, 0); - } - while (queue.size()) { - auto [token, next, prev] = queue.front(); - queue.pop(); - info.token.emplace_back(token); - prevs.emplace_back(prev); - for (auto [t, n] : tree[next].next) { - queue.emplace(t, n, info.token.size() - 1); - } - } - - // zero padding to length - while (info.token.size() < draft_token_num) { - info.token.emplace_back(0); - prevs.emplace_back(0); - } - - int n = info.token.size(); - info.mask.resize(n * n, 0); - info.mask[0] = 1; - for (int i = 0; i < n; ++i) { - if (prevs[i] != -1) { - memcpy(&info.mask[i * n], &info.mask[prevs[i] * n], prevs[i] + 1); - } - info.mask[i * n + i] = 1; - } - - return info; -} - -Ngram::Ngram(size_t capacity, const Param& param) { - param_ = param; - nodes_.resize(capacity); - for (auto& node : nodes_) { - node_pool_.emplace_back(&node); - } - free_node_count_ = node_pool_.size(); - root_ = getNode(); - +Ngram::Ngram(size_t capacity, const Param& param) : param_(param) { if (!(param_.branch_length > 1)) { throw std::runtime_error( "param_.branch_length must be greater than 1, current value: " + std::to_string(param_.branch_length)); @@ -79,13 +21,15 @@ Ngram::Ngram(size_t capacity, const Param& param) { } if (!(param_.min_match_window_size <= param_.max_match_window_size)) { throw std::runtime_error( - "min_match_window_size must be less than or equal to max_match_window_size, current min_match_window_size: " + + "min_match_window_size must be less than or equal to " + "max_match_window_size, current min_match_window_size: " + std::to_string(param_.min_match_window_size) + ", max_match_window_size: " + std::to_string(param_.max_match_window_size)); } if (!(param_.max_match_window_size < param_.branch_length)) { throw std::runtime_error( - "max_match_window_size must be less than branch_length, current max_match_window_size: " + + "max_match_window_size must be less than branch_length, current " + "max_match_window_size: " + std::to_string(param_.max_match_window_size) + ", branch_length: " + std::to_string(param_.branch_length)); } if (!(param_.min_bfs_breadth > 0)) { @@ -94,7 +38,8 @@ Ngram::Ngram(size_t capacity, const Param& param) { } if (!(param_.min_bfs_breadth <= param_.max_bfs_breadth)) { throw std::runtime_error( - "min_bfs_breadth must be less than or equal to max_bfs_breadth, current min_bfs_breadth: " + + "min_bfs_breadth must be less than or equal to max_bfs_breadth, " + "current min_bfs_breadth: " + std::to_string(param_.min_bfs_breadth) + ", max_bfs_breadth: " + std::to_string(param_.max_bfs_breadth)); } if (!(param_.draft_token_num > 0)) { @@ -125,64 +70,17 @@ Ngram::Ngram(size_t capacity, const Param& param) { } } + trie_ = std::make_unique(capacity, param_); + quit_flag_ = false; - insert_worker_ = std::thread(&Ngram::insert, this); + insert_worker_ = std::thread(&Ngram::insertWorker, this); } Ngram::~Ngram() { quit_flag_ = true; insert_queue_.close(); - insert_worker_.join(); -} - -std::vector> Ngram::match(const std::vector& tokens, size_t batch_size) const { - auto draft_token_num = param_.get_draft_token_num(batch_size); - auto min_match_window_size = param_.get_min_match_window_size(batch_size); - auto max_match_window_size = param_.max_match_window_size; - std::vector> result; - result.reserve(param_.max_match_window_size - param_.min_match_window_size); - for (int32_t match_window_size = std::min(tokens.size(), param_.max_match_window_size); - match_window_size >= param_.min_match_window_size; - --match_window_size) { - auto start = tokens.data() + tokens.size() - match_window_size; - auto end = start + match_window_size; - auto cursor = root_; - while (start != end) { - auto iter = cursor->child.find(*start); - if (iter == cursor->child.end()) { - cursor = nullptr; - break; - } - ++start; - cursor = iter->second; - } - if (cursor) { - result.emplace_back(std::make_pair(cursor, match_window_size)); - } - } - return result; -} - -void Ngram::squeeze(size_t count) { - if (!(node_pool_.size() >= free_node_count_ + count)) { - throw std::runtime_error( - "Insufficient node size to release required nodes. " - "available to release: " + - std::to_string(node_pool_.size() - free_node_count_) + ", required to release: " + std::to_string(count)); - } - while (count--) { - auto last = global_lru_.back(); - global_lru_.pop_back(); - - if (!last->child.empty()) { - throw std::runtime_error("The node to be released still has child nodes and cannot be released. "); - } - - last->parent->lru.erase(last->parent_lru_pos); - last->parent->sorted_children.erase(last); - last->parent->child.erase(last->token); - - node_pool_[free_node_count_++] = last; + if (insert_worker_.joinable()) { + insert_worker_.join(); } } @@ -192,190 +90,44 @@ void Ngram::synchronize() const { } } -void Ngram::insert() { - while (!quit_flag_) { - std::vector data; - if (!insert_queue_.dequeue(data)) { - continue; - } - const auto* token = data.data(); - size_t size = data.size(); - std::unique_lock lock(mutex_); - - for (size_t i = 0; i + param_.min_match_window_size < size; ++i) { - auto start = token + i; - auto end = start + std::min(size - i, param_.branch_length); - - if (end - start > free_node_count_) { - squeeze(end - start - free_node_count_); - } - - TrieNode* cursor = root_; - path_.clear(); - while (start != end) { - auto token = *start; - auto iter = cursor->child.find(token); - if (iter == cursor->child.end()) { - iter = cursor->child.insert({token, getNode()}).first; - auto node = iter->second; - - cursor->lru.emplace_front(node); - global_lru_.emplace_back(node); - - node->token = token; - node->parent = cursor; - node->parent_lru_pos = cursor->lru.begin(); - node->global_lru_pos = --global_lru_.end(); - node->freq = 1; - cursor->sorted_children.insert(node); - } else { - auto node = iter->second; - cursor->sorted_children.erase(node); - node->freq++; - cursor->sorted_children.insert(node); - cursor->lru.splice(cursor->lru.begin(), cursor->lru, node->parent_lru_pos); - } - cursor = iter->second; - path_.emplace_back(cursor); - ++start; - } - - for (auto it = path_.rbegin(); it != path_.rend(); ++it) { - TrieNode* node = *it; - global_lru_.splice(global_lru_.begin(), global_lru_, node->global_lru_pos); - } - } - } -} - void Ngram::asyncInsert(std::vector>&& tokens) { for (auto&& token : tokens) { insert_queue_.enqueue(std::move(token)); } } -Ngram::Result Ngram::matchBFS(const std::vector& tokens, size_t batch_size) const { - std::vector> nodes = match(tokens, batch_size); - - double bfs_breadth_scale = double(param_.max_bfs_breadth - param_.min_bfs_breadth) / - (param_.max_match_window_size - param_.min_match_window_size + 1); - - auto draft_token_num = param_.get_draft_token_num(batch_size); - std::vector tree(draft_token_num + 1); - int root = 0; - int cursor = 1; - - for (auto [node, depth] : nodes) { - std::queue> queue; // parent, bfs_breadth, node - queue.push({root, (param_.max_match_window_size - depth) * bfs_breadth_scale + param_.min_bfs_breadth, node}); - while (queue.size() && cursor <= draft_token_num) { - auto front = queue.front(); - queue.pop(); - - auto parent = std::get<0>(front); - auto cur_breadth = std::get<1>(front); - auto iter = std::get<2>(front)->lru.begin(); - - auto breadth = std::max(1, int32_t(cur_breadth)); - for (int i = 0; i < breadth && iter != std::get<2>(front)->lru.end() && cursor <= draft_token_num; ++i, ++iter) { - auto token = (*iter)->token; - auto pos = -1; - if (auto tit = tree[parent].next.find(token); tit != tree[parent].next.end()) { - pos = tit->second; - } else { - pos = tree[parent].next.insert(std::make_pair(token, cursor++)).first->second; - } - queue.emplace(pos, cur_breadth - bfs_breadth_scale, *iter); - } +void Ngram::insertWorker() { + while (!quit_flag_) { + std::vector data; + if (!insert_queue_.dequeue(data)) { + continue; } + std::unique_lock lock(mutex_); + trie_->insert(data.data(), data.size()); } - - return fillResult(tokens.back(), draft_token_num + 1, tree, root); } -Ngram::Result Ngram::matchProb(const std::vector& tokens, size_t batch_size) const { - std::vector> nodes = match(tokens, batch_size); - auto draft_token_num = param_.get_draft_token_num(batch_size); - - struct CompareByLastDouble { - bool operator()( - const std::tuple& a, // parent_pos, node, final_prob - const std::tuple& b) const { - return std::get<2>(a) < std::get<2>(b); - } - }; - - std::priority_queue< - std::tuple, - std::vector>, - CompareByLastDouble> - heap; - - std::vector tree(draft_token_num + 1); - - int root = 0; - int cursor = 1; - int top_k = param_.max_bfs_breadth; - - auto addToHeap = [&heap, &top_k](int parent, const TrieNode* trie_node, double prob) -> void { - double sum_freq = 0.0; - int count = 0; - std::list> topk_children; - for (auto* child : trie_node->sorted_children) { - sum_freq += static_cast(child->freq); - topk_children.emplace_back(child, child->freq); - if (++count >= top_k) break; - } - if (sum_freq <= 0) sum_freq = 1.0; - for (const auto& [child, freq] : topk_children) { - double norm_freq = static_cast(freq) / sum_freq * prob; - heap.emplace(parent, child, norm_freq); - } - }; - - for (auto [node, _] : nodes) { - addToHeap(root, node, 1.0); - - while (!heap.empty() && cursor <= draft_token_num) { - auto [parent, trie_node, prob] = heap.top(); // parent_pos, node, final_prob - heap.pop(); - auto token = trie_node->token; - int pos = -1; - auto tit = tree[parent].next.find(token); - if (tit != tree[parent].next.end()) { - pos = tit->second; - } else { - pos = cursor++; - tree[parent].next[token] = pos; - } - addToHeap(pos, trie_node, prob); - } - } - - return fillResult(tokens.back(), draft_token_num + 1, tree, root); -} - -Ngram::Result Ngram::batchMatch(const std::vector>& tokens) const { +Result Ngram::batchMatch(const std::vector>& tokens) const { std::unique_lock lock(mutex_); - Result merged_result; - auto match_func = param_.match_type == "BFS" ? &Ngram::matchBFS : &Ngram::matchProb; - for (const auto& tks : tokens) { - Result res = (this->*match_func)(tks, tokens.size()); - merged_result.token.insert(merged_result.token.end(), res.token.begin(), res.token.end()); - merged_result.mask.insert(merged_result.mask.end(), res.mask.begin(), res.mask.end()); - } - return merged_result; -} -void Ngram::Result::truncate(size_t n) { - if (n < token.size()) { - int full_n = token.size(); - for (int i = 1; i < n; ++i) { - memcpy(&mask[i * n], &mask[i * full_n], sizeof(mask[0]) * n); - } - token.resize(n); - mask.resize(n * n); + using BuildFn = Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&) const; + BuildFn build_fn; + if (param_.match_type == "BFS") { + build_fn = &Trie::buildRecency; + } else if (param_.match_type == "PROB") { + build_fn = &Trie::buildFrequency; + } else { + throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'."); } + + Result merged; + for (const auto& suffix : tokens) { + auto draft_token_num = param_.get_draft_token_num(tokens.size()); + auto res = (trie_.get()->*build_fn)(suffix.data(), suffix.size(), suffix.back(), draft_token_num, param_); + merged.token.insert(merged.token.end(), res.token.begin(), res.token.end()); + merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end()); + } + return merged; } } // namespace ngram diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.h b/python/sglang/srt/speculative/cpp_ngram/ngram.h index 3c9a9380e..9af65885b 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram.h +++ b/python/sglang/srt/speculative/cpp_ngram/ngram.h @@ -2,99 +2,42 @@ #include #include -#include -#include +#include #include -#include -#include -#include #include -#include -#include #include #include "param.h" #include "queue.h" +#include "result.h" +#include "trie.h" namespace ngram { -struct TrieNode { - std::unordered_map child; - std::list::const_iterator global_lru_pos; - std::list::const_iterator parent_lru_pos; - int32_t token; - TrieNode* parent; - std::list lru; - int32_t freq = 0; - - struct CompareByFreq { - bool operator()(TrieNode* a, TrieNode* b) const { - return std::tie(b->freq, a->token, a) < std::tie(a->freq, b->token, b); - } - }; - std::multiset sorted_children; -}; - class Ngram { - std::vector nodes_; - std::vector node_pool_; - size_t free_node_count_; - std::list global_lru_; - TrieNode* root_; - std::vector path_; + std::unique_ptr trie_; Param param_; - std::vector> match(const std::vector& tokens, size_t batch_size) const; - - void squeeze(size_t count); - - TrieNode* getNode() { - auto node = node_pool_[--free_node_count_]; - node->~TrieNode(); - new (node) TrieNode(); - return node; - } - mutable std::mutex mutex_; - bool quit_flag_; + bool quit_flag_ = false; utils::Queue> insert_queue_; std::thread insert_worker_; - std::vector> match_tmp_data_; public: Ngram(size_t capacity, const Param& param); - Ngram() = default; ~Ngram(); - static Ngram& instance() { - static Ngram instance; - return instance; - } - void synchronize() const; void asyncInsert(std::vector>&& tokens); - struct Result { - std::vector token; - std::vector mask; - - void truncate(size_t n); - }; - Result batchMatch(const std::vector>& tokens) const; void reset() { std::unique_lock lock(mutex_); - - global_lru_.clear(); - path_.clear(); - node_pool_.clear(); - for (auto& node : nodes_) { - node_pool_.emplace_back(&node); + if (trie_) { + trie_->reset(); } - free_node_count_ = node_pool_.size(); - root_ = getNode(); } const Param& param() const { @@ -102,10 +45,7 @@ class Ngram { } private: - Result matchBFS(const std::vector& tokens, size_t batch_size) const; - Result matchProb(const std::vector& tokens, size_t batch_size) const; - - void insert(); + void insertWorker(); }; } // namespace ngram diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_cache.py b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py similarity index 83% rename from python/sglang/srt/speculative/cpp_ngram/ngram_cache.py rename to python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py index 8b1eb8eea..56d417bd4 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram_cache.py +++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py @@ -10,17 +10,19 @@ from torch.utils.cpp_extension import load logger = logging.getLogger(__name__) _abs_path = os.path.dirname(os.path.abspath(__file__)) -ngram_cache_cpp = load( - name="ngram_cache_cpp", +ngram_corpus_cpp = load( + name="ngram_corpus_cpp", sources=[ - f"{_abs_path}/ngram_cache_binding.cpp", + f"{_abs_path}/ngram_corpus_binding.cpp", f"{_abs_path}/ngram.cpp", + f"{_abs_path}/trie.cpp", + f"{_abs_path}/result.cpp", ], extra_cflags=["-O3", "-std=c++20"], ) -class NgramCache: +class NgramCorpus: def __init__( self, branch_length=18, @@ -32,7 +34,7 @@ class NgramCache: match_type="BFS", capacity=1000000, ): - param = ngram_cache_cpp.Param() + param = ngram_corpus_cpp.Param() param.branch_length = branch_length param.min_match_window_size = min_match_window_size param.max_match_window_size = max_match_window_size @@ -40,22 +42,22 @@ class NgramCache: param.max_bfs_breadth = max_bfs_breadth param.draft_token_num = draft_token_num param.match_type = match_type - self.cache = ngram_cache_cpp.Ngram(capacity, param) + self._ngram = ngram_corpus_cpp.Ngram(capacity, param) self.default_mask = np.ones((1, 1), dtype=np.int64) self.draft_token_num = draft_token_num def batch_put(self, batch_tokens: List[List[int]]): - self.cache.asyncInsert(batch_tokens) + self._ngram.asyncInsert(batch_tokens) def synchronize(self): - self.cache.synchronize() + self._ngram.synchronize() def reset(self): - self.cache.reset() + self._ngram.reset() def batch_get(self, batch_tokens: List[List[int]]) -> Tuple[np.ndarray, np.ndarray]: - result = self.cache.batchMatch(batch_tokens) + result = self._ngram.batchMatch(batch_tokens) return np.array(result.token), np.array(result.mask) def leaf_paths_from_mask( @@ -129,10 +131,10 @@ if __name__ == "__main__": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 44, 55, 66, 77, 88, 99, 100], ] - cache = NgramCache(branch_length=12, draft_token_num=8) - cache.batch_put(token_ids) + corpus = NgramCorpus(branch_length=12, draft_token_num=8) + corpus.batch_put(token_ids) - cache.synchronize() - decoding_ids, decoding_masks = cache.batch_get([[1, 2, 3], [3, 44], [3, 6, 999]]) + corpus.synchronize() + decoding_ids, decoding_masks = corpus.batch_get([[1, 2, 3], [3, 44], [3, 6, 999]]) - cache.debug_result(decoding_ids, decoding_masks) + corpus.debug_result(decoding_ids, decoding_masks) diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp similarity index 87% rename from python/sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp rename to python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp index ac5b931f9..edc0e2c5c 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp +++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp @@ -3,7 +3,7 @@ #include "ngram.h" -PYBIND11_MODULE(ngram_cache_cpp, m) { +PYBIND11_MODULE(ngram_corpus_cpp, m) { using namespace ngram; namespace py = pybind11; m.doc() = ""; @@ -35,9 +35,9 @@ PYBIND11_MODULE(ngram_cache_cpp, m) { .def("resetBatchReturnTokenNum", &Param::resetBatchReturnTokenNum, "") .def("detail", &Param::detail, ""); - py::class_(m, "Result") + py::class_(m, "Result") .def(py::init<>()) - .def_readwrite("token", &Ngram::Result::token) - .def_readwrite("mask", &Ngram::Result::mask) - .def("truncate", &Ngram::Result::truncate); + .def_readwrite("token", &Result::token) + .def_readwrite("mask", &Result::mask) + .def("truncate", &Result::truncate); } diff --git a/python/sglang/srt/speculative/cpp_ngram/result.cpp b/python/sglang/srt/speculative/cpp_ngram/result.cpp new file mode 100644 index 000000000..404b7a3f2 --- /dev/null +++ b/python/sglang/srt/speculative/cpp_ngram/result.cpp @@ -0,0 +1,61 @@ +#include "result.h" + +#include +#include +#include + +namespace ngram { + +Result fillResult(int last_token, int draft_token_num, std::vector& tree, int root) { + Result info; + std::vector prevs; + info.token.reserve(draft_token_num); + prevs.reserve(draft_token_num); + std::queue> queue; + info.token.emplace_back(last_token); + prevs.emplace_back(-1); + + for (auto [token, next] : tree[root].next) { + queue.emplace(token, next, 0); + } + while (queue.size()) { + auto [token, next, prev] = queue.front(); + queue.pop(); + info.token.emplace_back(token); + prevs.emplace_back(prev); + for (auto [t, n] : tree[next].next) { + queue.emplace(t, n, info.token.size() - 1); + } + } + + // zero padding to length + while (info.token.size() < static_cast(draft_token_num)) { + info.token.emplace_back(0); + prevs.emplace_back(0); + } + + int n = info.token.size(); + info.mask.resize(n * n, 0); + info.mask[0] = 1; + for (int i = 0; i < n; ++i) { + if (prevs[i] != -1) { + memcpy(&info.mask[i * n], &info.mask[prevs[i] * n], prevs[i] + 1); + } + info.mask[i * n + i] = 1; + } + + return info; +} + +void Result::truncate(size_t n) { + if (n < token.size()) { + int full_n = token.size(); + for (size_t i = 1; i < n; ++i) { + memcpy(&mask[i * n], &mask[i * full_n], sizeof(mask[0]) * n); + } + token.resize(n); + mask.resize(n * n); + } +} + +} // namespace ngram diff --git a/python/sglang/srt/speculative/cpp_ngram/result.h b/python/sglang/srt/speculative/cpp_ngram/result.h new file mode 100644 index 000000000..c48351d77 --- /dev/null +++ b/python/sglang/srt/speculative/cpp_ngram/result.h @@ -0,0 +1,22 @@ +#pragma once + +#include +#include +#include + +namespace ngram { + +struct Result { + std::vector token; + std::vector mask; + + void truncate(size_t n); +}; + +struct Node { + std::unordered_map next; +}; + +Result fillResult(int last_token, int draft_token_num, std::vector& tree, int root); + +} // namespace ngram diff --git a/python/sglang/srt/speculative/cpp_ngram/trie.cpp b/python/sglang/srt/speculative/cpp_ngram/trie.cpp new file mode 100644 index 000000000..7c7340d57 --- /dev/null +++ b/python/sglang/srt/speculative/cpp_ngram/trie.cpp @@ -0,0 +1,231 @@ +#include "trie.h" + +#include +#include +#include +#include +#include +#include + +namespace ngram { + +Trie::Trie(size_t capacity, const Param& param) : param_(param) { + nodes_.resize(capacity); + for (auto& node : nodes_) { + node_pool_.emplace_back(&node); + } + free_node_count_ = node_pool_.size(); + root_ = getNode(); +} + +void Trie::insert(const int32_t* tokens, size_t len) { + for (size_t i = 0; i + param_.min_match_window_size < len; ++i) { + auto start = tokens + i; + auto end = start + std::min(len - i, param_.branch_length); + + if (static_cast(end - start) > free_node_count_) { + squeeze(end - start - free_node_count_); + } + + TrieNode* cursor = root_; + path_.clear(); + while (start != end) { + auto token = *start; + auto iter = cursor->child.find(token); + if (iter == cursor->child.end()) { + iter = cursor->child.insert({token, getNode()}).first; + auto node = iter->second; + + cursor->lru.emplace_front(node); + global_lru_.emplace_back(node); + + node->token = token; + node->parent = cursor; + node->parent_lru_pos = cursor->lru.begin(); + node->global_lru_pos = --global_lru_.end(); + node->freq = 1; + cursor->sorted_children.insert(node); + } else { + auto node = iter->second; + cursor->sorted_children.erase(node); + node->freq++; + cursor->sorted_children.insert(node); + cursor->lru.splice(cursor->lru.begin(), cursor->lru, node->parent_lru_pos); + } + cursor = iter->second; + path_.emplace_back(cursor); + ++start; + } + + for (auto it = path_.rbegin(); it != path_.rend(); ++it) { + TrieNode* node = *it; + global_lru_.splice(global_lru_.begin(), global_lru_, node->global_lru_pos); + } + } +} + +void Trie::squeeze(size_t count) { + if (!(node_pool_.size() >= free_node_count_ + count)) { + throw std::runtime_error( + "Insufficient node size to release required nodes. " + "available to release: " + + std::to_string(node_pool_.size() - free_node_count_) + ", required to release: " + std::to_string(count)); + } + while (count--) { + auto last = global_lru_.back(); + global_lru_.pop_back(); + + if (!last->child.empty()) { + throw std::runtime_error( + "The node to be released still has child nodes and cannot be " + "released. "); + } + + last->parent->lru.erase(last->parent_lru_pos); + last->parent->sorted_children.erase(last); + last->parent->child.erase(last->token); + + node_pool_[free_node_count_++] = last; + } +} + +void Trie::reset() { + global_lru_.clear(); + path_.clear(); + node_pool_.clear(); + for (auto& node : nodes_) { + node_pool_.emplace_back(&node); + } + free_node_count_ = node_pool_.size(); + root_ = getNode(); +} + +std::vector> +Trie::match(const int32_t* context, size_t len, size_t min_window, size_t max_window) const { + std::vector> result; + result.reserve(max_window - min_window); + for (int32_t match_window_size = std::min(len, max_window); match_window_size >= static_cast(min_window); + --match_window_size) { + auto start = context + len - match_window_size; + auto end = start + match_window_size; + auto cursor = root_; + while (start != end) { + auto iter = cursor->child.find(*start); + if (iter == cursor->child.end()) { + cursor = nullptr; + break; + } + ++start; + cursor = iter->second; + } + if (cursor) { + result.emplace_back(std::make_pair(cursor, match_window_size)); + } + } + return result; +} + +Result Trie::buildRecency( + const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const { + auto anchors = match(context, len, param.min_match_window_size, param.max_match_window_size); + + double bfs_breadth_scale = double(param.max_bfs_breadth - param.min_bfs_breadth) / + (param.max_match_window_size - param.min_match_window_size + 1); + + std::vector tree(draft_token_num + 1); + int root = 0; + int cursor = 1; + + for (auto [node, depth] : anchors) { + std::queue> queue; + queue.push({root, (param.max_match_window_size - depth) * bfs_breadth_scale + param.min_bfs_breadth, node}); + while (queue.size() && cursor <= static_cast(draft_token_num)) { + auto front = queue.front(); + queue.pop(); + + auto parent = std::get<0>(front); + auto cur_breadth = std::get<1>(front); + auto iter = std::get<2>(front)->lru.begin(); + + auto breadth = std::max(1, int32_t(cur_breadth)); + for (int i = 0; + i < breadth && iter != std::get<2>(front)->lru.end() && cursor <= static_cast(draft_token_num); + ++i, ++iter) { + auto token = (*iter)->token; + auto pos = -1; + if (auto tit = tree[parent].next.find(token); tit != tree[parent].next.end()) { + pos = tit->second; + } else { + pos = tree[parent].next.insert(std::make_pair(token, cursor++)).first->second; + } + queue.emplace(pos, cur_breadth - bfs_breadth_scale, *iter); + } + } + } + + return fillResult(last_token, draft_token_num + 1, tree, root); +} + +Result Trie::buildFrequency( + const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const { + auto anchors = match(context, len, param.min_match_window_size, param.max_match_window_size); + + struct CompareByLastDouble { + bool operator()( + const std::tuple& a, + const std::tuple& b) const { + return std::get<2>(a) < std::get<2>(b); + } + }; + + std::priority_queue< + std::tuple, + std::vector>, + CompareByLastDouble> + heap; + + std::vector tree(draft_token_num + 1); + + int root = 0; + int cursor = 1; + int top_k = param.max_bfs_breadth; + + auto addToHeap = [&heap, &top_k](int parent, const TrieNode* trie_node, double prob) -> void { + double sum_freq = 0.0; + int count = 0; + std::list> topk_children; + for (auto* child : trie_node->sorted_children) { + sum_freq += static_cast(child->freq); + topk_children.emplace_back(child, child->freq); + if (++count >= top_k) break; + } + if (sum_freq <= 0) sum_freq = 1.0; + for (const auto& [child, freq] : topk_children) { + double norm_freq = static_cast(freq) / sum_freq * prob; + heap.emplace(parent, child, norm_freq); + } + }; + + for (auto [node, _] : anchors) { + addToHeap(root, node, 1.0); + + while (!heap.empty() && cursor <= static_cast(draft_token_num)) { + auto [parent, trie_node, prob] = heap.top(); + heap.pop(); + auto token = trie_node->token; + int pos = -1; + auto tit = tree[parent].next.find(token); + if (tit != tree[parent].next.end()) { + pos = tit->second; + } else { + pos = cursor++; + tree[parent].next[token] = pos; + } + addToHeap(pos, trie_node, prob); + } + } + + return fillResult(last_token, draft_token_num + 1, tree, root); +} + +} // namespace ngram diff --git a/python/sglang/srt/speculative/cpp_ngram/trie.h b/python/sglang/srt/speculative/cpp_ngram/trie.h new file mode 100644 index 000000000..30db5b294 --- /dev/null +++ b/python/sglang/srt/speculative/cpp_ngram/trie.h @@ -0,0 +1,71 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "param.h" +#include "result.h" + +namespace ngram { + +struct TrieNode { + std::unordered_map child; + std::list::const_iterator global_lru_pos; + std::list::const_iterator parent_lru_pos; + int32_t token; + TrieNode* parent; + std::list lru; + int32_t freq = 0; + + struct CompareByFreq { + bool operator()(TrieNode* a, TrieNode* b) const { + return std::tie(b->freq, a->token, a) < std::tie(a->freq, b->token, b); + } + }; + std::multiset sorted_children; +}; + +class Trie { + public: + Trie(size_t capacity, const Param& param); + + void insert(const int32_t* tokens, size_t len); + + Result buildRecency( + const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const; + + Result buildFrequency( + const int32_t* context, size_t len, int32_t last_token, size_t draft_token_num, const Param& param) const; + + void squeeze(size_t count); + + void reset(); + + private: + std::vector> + match(const int32_t* context, size_t len, size_t min_window, size_t max_window) const; + + TrieNode* getNode() { + auto node = node_pool_[--free_node_count_]; + node->~TrieNode(); + new (node) TrieNode(); + return node; + } + + std::vector nodes_; + std::vector node_pool_; + size_t free_node_count_; + std::list global_lru_; + TrieNode* root_; + std::vector path_; + Param param_; +}; + +} // namespace ngram diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 7f6277bb8..784a4130e 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -11,7 +11,7 @@ from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.server_args import ServerArgs -from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache +from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus from sglang.srt.speculative.ngram_info import NgramVerifyInput from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import generate_token_bitmask @@ -50,18 +50,19 @@ class NGRAMWorker: self._init_preallocated_tensors() - self.ngram_cache = NgramCache( + self.ngram_corpus = NgramCorpus( min_match_window_size=server_args.speculative_ngram_min_match_window_size, max_match_window_size=server_args.speculative_ngram_max_match_window_size, min_bfs_breadth=server_args.speculative_ngram_min_bfs_breadth, max_bfs_breadth=server_args.speculative_ngram_max_bfs_breadth, + match_type=server_args.speculative_ngram_match_type, capacity=server_args.speculative_ngram_capacity, branch_length=server_args.speculative_ngram_branch_length, draft_token_num=server_args.speculative_num_draft_tokens, ) def clear_cache_pool(self): - self.ngram_cache.reset() + self.ngram_corpus.reset() def _efficient_concat_last_n(self, seq1: List[int], seq2: List[int], n: int): seq2_len = len(seq2) @@ -126,14 +127,14 @@ class NGRAMWorker: ) -> tuple[np.ndarray, np.ndarray]: bs = batch.batch_size() - self.ngram_cache.synchronize() + self.ngram_corpus.synchronize() batch_tokens = [] for req in batch.reqs: check_token = self._efficient_concat_last_n( req.origin_input_ids, req.output_ids, self.max_match_window_size ) batch_tokens.append(check_token) - req_drafts, mask = self.ngram_cache.batch_get(batch_tokens) + req_drafts, mask = self.ngram_corpus.batch_get(batch_tokens) total_draft_token_num = len(req_drafts) # Check if speculative decoding is needed; here we always enforce it @@ -199,7 +200,7 @@ class NGRAMWorker: ) batch.spec_info.prepare_for_verify(batch, self.page_size) - def _update_ngram_cache(self, batch: ScheduleBatch): + def _update_ngram_corpus(self, batch: ScheduleBatch): batch_tokens = [] for req in batch.reqs: # FIXME: Whether to insert 'extend' into the cache or not, after testing, @@ -211,7 +212,7 @@ class NGRAMWorker: req.origin_input_ids, req.output_ids, self.branch_length ) batch_tokens.append(put_ids) - self.ngram_cache.batch_put(batch_tokens) + self.ngram_corpus.batch_put(batch_tokens) def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult: self._prepare_for_speculative_decoding(batch) @@ -264,7 +265,7 @@ class NGRAMWorker: accept_lens = verify_input.accept_length if batch.return_logprob: add_output_logprobs_for_spec_v1(batch, verify_input, logits_output) - self._update_ngram_cache(batch) + self._update_ngram_corpus(batch) batch.forward_mode = ForwardMode.DECODE else: diff --git a/test/registered/spec/utils/test_ngram_corpus.py b/test/registered/spec/utils/test_ngram_corpus.py new file mode 100644 index 000000000..de1b20168 --- /dev/null +++ b/test/registered/spec/utils/test_ngram_corpus.py @@ -0,0 +1,571 @@ +import unittest + +import numpy as np + +from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, suite="stage-b-test-small-1-gpu") + + +def _make_corpus(match_type="BFS", **kwargs): + defaults = dict( + branch_length=12, + min_match_window_size=1, + max_match_window_size=10, + min_bfs_breadth=1, + max_bfs_breadth=8, + draft_token_num=8, + capacity=100000, + ) + defaults.update(kwargs) + defaults["match_type"] = match_type + return NgramCorpus(**defaults) + + +SEED_SEQUENCES = [ + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + [1, 2, 3, 44, 55, 66, 77, 88, 99, 100], +] + +QUERY_SEQUENCES = [[1, 2, 3], [3, 44], [3, 6, 999]] + +EXPECTED_BFS_IDS = [ + [3, 4, 44, 5, 55, 6, 66, 77], + [44, 55, 66, 77, 88, 99, 100, 0], + [999, 0, 0, 0, 0, 0, 0, 0], +] + +EXPECTED_PROB_IDS = [ + [3, 44, 4, 55, 5, 66, 6, 7], + [44, 55, 66, 77, 88, 99, 100, 0], + [999, 0, 0, 0, 0, 0, 0, 0], +] + +EXPECTED_BFS_MASKS = [ + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0, 0, 0], + [1, 1, 0, 1, 0, 0, 0, 0], + [1, 0, 1, 0, 1, 0, 0, 0], + [1, 1, 0, 1, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 0, 1, 0], + [1, 0, 1, 0, 1, 0, 1, 1], + ], + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 1, 0, 0], + [1, 1, 1, 1, 1, 1, 1, 0], + [1, 0, 0, 0, 0, 0, 0, 1], + ], + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 0, 0, 0, 0], + [1, 0, 0, 0, 1, 0, 0, 0], + [1, 0, 0, 0, 0, 1, 0, 0], + [1, 0, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 0, 0, 0, 0, 1], + ], +] + +EXPECTED_PROB_MASKS = [ + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0, 0, 0], + [1, 1, 0, 1, 0, 0, 0, 0], + [1, 0, 1, 0, 1, 0, 0, 0], + [1, 1, 0, 1, 0, 1, 0, 0], + [1, 0, 1, 0, 1, 0, 1, 0], + [1, 0, 1, 0, 1, 0, 1, 1], + ], + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 1, 1, 0, 0, 0, 0, 0], + [1, 1, 1, 1, 0, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 1, 0, 0], + [1, 1, 1, 1, 1, 1, 1, 0], + [1, 0, 0, 0, 0, 0, 0, 1], + ], + [ + [1, 0, 0, 0, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0, 0, 0], + [1, 0, 1, 0, 0, 0, 0, 0], + [1, 0, 0, 1, 0, 0, 0, 0], + [1, 0, 0, 0, 1, 0, 0, 0], + [1, 0, 0, 0, 0, 1, 0, 0], + [1, 0, 0, 0, 0, 0, 1, 0], + [1, 0, 0, 0, 0, 0, 0, 1], + ], +] + + +class TestNgramCorpusBFS(CustomTestCase): + """Golden-output tests for BFS matching mode.""" + + @classmethod + def setUpClass(cls): + cls.corpus = _make_corpus("BFS") + cls.corpus.batch_put(SEED_SEQUENCES) + cls.corpus.synchronize() + ids, masks = cls.corpus.batch_get(QUERY_SEQUENCES) + draft = 8 + cls.ids = ids.reshape(-1, draft) + cls.masks = masks.reshape(-1, draft, draft) + + def test_token_ids(self): + np.testing.assert_array_equal(self.ids.tolist(), EXPECTED_BFS_IDS) + + def test_masks(self): + np.testing.assert_array_equal(self.masks.tolist(), EXPECTED_BFS_MASKS) + + def test_output_shapes(self): + n_queries = len(QUERY_SEQUENCES) + draft = 8 + self.assertEqual(self.ids.shape, (n_queries, draft)) + self.assertEqual(self.masks.shape, (n_queries, draft, draft)) + + +class TestNgramCorpusProb(CustomTestCase): + """Golden-output tests for Prob matching mode.""" + + @classmethod + def setUpClass(cls): + cls.corpus = _make_corpus("PROB") + cls.corpus.batch_put(SEED_SEQUENCES) + cls.corpus.synchronize() + ids, masks = cls.corpus.batch_get(QUERY_SEQUENCES) + cls.ids = ids.reshape(-1, 8) + cls.masks = masks.reshape(-1, 8, 8) + + def test_token_ids(self): + np.testing.assert_array_equal(self.ids.tolist(), EXPECTED_PROB_IDS) + + def test_masks(self): + np.testing.assert_array_equal(self.masks.tolist(), EXPECTED_PROB_MASKS) + + def test_output_shapes(self): + n_queries = len(QUERY_SEQUENCES) + self.assertEqual(self.ids.shape, (n_queries, 8)) + self.assertEqual(self.masks.shape, (n_queries, 8, 8)) + + +class TestNgramCorpusReset(CustomTestCase): + """Verify reset clears all cached state.""" + + def test_reset_produces_empty_results(self): + corpus = _make_corpus("BFS") + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + + ids_before, _ = corpus.batch_get([[1, 2, 3]]) + self.assertTrue( + any(t != 0 for t in ids_before.tolist()[1:]), + "Expected non-trivial draft tokens before reset", + ) + + corpus.reset() + + ids_after, _ = corpus.batch_get([[1, 2, 3]]) + self.assertEqual( + ids_after.tolist(), + [3, 0, 0, 0, 0, 0, 0, 0], + "After reset, only last_token should be present (rest zero-padded)", + ) + + +class TestNgramCorpusNoMatch(CustomTestCase): + """Verify behavior when query has no match in the corpus.""" + + def test_unmatched_query(self): + corpus = _make_corpus("BFS") + corpus.batch_put([[10, 20, 30, 40, 50]]) + corpus.synchronize() + + ids, masks = corpus.batch_get([[999, 888, 777]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 777, "First token should be last context token") + self.assertTrue( + all(t == 0 for t in ids_list[1:]), + "No draft tokens expected when nothing matches", + ) + + def test_empty_corpus(self): + corpus = _make_corpus("BFS") + ids, masks = corpus.batch_get([[1, 2, 3]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 3) + self.assertTrue(all(t == 0 for t in ids_list[1:])) + + +class TestNgramCorpusMultipleInserts(CustomTestCase): + """Verify that multiple inserts accumulate correctly.""" + + def test_incremental_inserts(self): + corpus = _make_corpus("BFS") + corpus.batch_put([[1, 2, 3, 4, 5]]) + corpus.synchronize() + + corpus.batch_put([[1, 2, 3, 44, 55]]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[1, 2, 3]]) + ids_list = ids.tolist() + + self.assertIn(4, ids_list, "Token 4 from first insert should still match") + self.assertIn(44, ids_list, "Token 44 from second insert should also match") + + +class TestNgramCorpusSqueeze(CustomTestCase): + """Verify cache eviction under memory pressure.""" + + def test_small_capacity_does_not_crash(self): + corpus = _make_corpus("BFS", capacity=200) + long_seq = list(range(1, 101)) + corpus.batch_put([long_seq]) + corpus.synchronize() + + ids, masks = corpus.batch_get([[50, 51, 52]]) + self.assertEqual(len(ids), 8, "Should still produce draft_token_num outputs") + + def test_eviction_preserves_recent(self): + corpus = _make_corpus( + "BFS", capacity=500, branch_length=6, max_match_window_size=5 + ) + + old_seq = list(range(1000, 1050)) + corpus.batch_put([old_seq]) + corpus.synchronize() + + recent_seq = list(range(2000, 2050)) + corpus.batch_put([recent_seq]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[2000, 2001, 2002]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 2002, "Last context token should be first") + self.assertIn(2003, ids_list, "Recent sequence should still be matchable") + + +class TestNgramCorpusLeafPaths(CustomTestCase): + """Verify the leaf_paths_from_mask utility.""" + + def test_simple_tree(self): + corpus = _make_corpus("BFS") + tokens = [3, 4, 44, 5, 55] + mask = [ + [1, 0, 0, 0, 0], + [1, 1, 0, 0, 0], + [1, 0, 1, 0, 0], + [1, 1, 0, 1, 0], + [1, 0, 1, 0, 1], + ] + paths = corpus.leaf_paths_from_mask(tokens, mask) + + for path in paths: + self.assertIn(3, path, "Root token should be in every path") + + self.assertEqual(len(paths), 2, "Two leaf paths expected for a binary tree") + + def test_single_chain(self): + corpus = _make_corpus("BFS") + tokens = [10, 20, 30] + mask = [ + [1, 0, 0], + [1, 1, 0], + [1, 1, 1], + ] + paths = corpus.leaf_paths_from_mask(tokens, mask) + self.assertEqual(len(paths), 1) + self.assertEqual(paths[0], [10, 20, 30]) + + +class TestNgramCorpusBatchConsistency(CustomTestCase): + """Verify batch queries produce same results as individual queries.""" + + def test_batch_vs_individual(self): + corpus = _make_corpus("BFS") + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + + batch_ids, batch_masks = corpus.batch_get(QUERY_SEQUENCES) + draft = 8 + batch_ids = batch_ids.reshape(-1, draft) + batch_masks = batch_masks.reshape(-1, draft, draft) + + for i, query in enumerate(QUERY_SEQUENCES): + single_ids, single_masks = corpus.batch_get([query]) + single_ids = single_ids.reshape(-1, draft) + single_masks = single_masks.reshape(-1, draft, draft) + + np.testing.assert_array_equal( + batch_ids[i], + single_ids[0], + err_msg=f"Token mismatch for query {i}", + ) + np.testing.assert_array_equal( + batch_masks[i], + single_masks[0], + err_msg=f"Mask mismatch for query {i}", + ) + + +class TestMaskValidity(CustomTestCase): + """Verify structural invariants of the output mask for any draft tree.""" + + def _check_mask(self, masks_2d): + n = len(masks_2d) + for i in range(n): + self.assertEqual(masks_2d[i][i], 1, f"Diagonal must be 1 at row {i}") + self.assertEqual(masks_2d[0], [1] + [0] * (n - 1)) + + def test_bfs_mask_invariants(self): + corpus = _make_corpus("BFS") + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + _, masks = corpus.batch_get(QUERY_SEQUENCES) + masks = masks.reshape(-1, 8, 8) + for i in range(masks.shape[0]): + self._check_mask(masks[i].tolist()) + + def test_prob_mask_invariants(self): + corpus = _make_corpus("PROB") + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + _, masks = corpus.batch_get(QUERY_SEQUENCES) + masks = masks.reshape(-1, 8, 8) + for i in range(masks.shape[0]): + self._check_mask(masks[i].tolist()) + + +class TestFrequencyBoosting(CustomTestCase): + """Verify that repeated insertions change Prob-mode selection.""" + + def test_repeated_insert_promotes_token(self): + corpus = _make_corpus( + "PROB", + draft_token_num=2, + max_bfs_breadth=1, + min_bfs_breadth=1, + max_match_window_size=3, + branch_length=5, + ) + corpus.batch_put([[1, 2, 3, 10, 11]]) + corpus.synchronize() + + for _ in range(10): + corpus.batch_put([[1, 2, 3, 20, 21]]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[1, 2, 3]]) + ids_list = ids.tolist() + + self.assertEqual( + ids_list[1], + 20, + f"Token 20 should be selected over 10 after frequency boost, got {ids_list}", + ) + + +class TestRecencyOrdering(CustomTestCase): + """Verify that BFS mode respects LRU recency.""" + + def test_most_recent_insert_selected(self): + corpus = _make_corpus( + "BFS", + draft_token_num=2, + max_bfs_breadth=1, + min_bfs_breadth=1, + max_match_window_size=3, + branch_length=5, + ) + corpus.batch_put([[1, 2, 3, 10, 11]]) + corpus.synchronize() + corpus.batch_put([[1, 2, 3, 20, 21]]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[1, 2, 3]]) + ids_list = ids.tolist() + self.assertEqual( + ids_list[1], + 20, + f"Token 20 (recent) should be selected over 10 (old), got {ids_list}", + ) + + +class TestOverlappingSuffixes(CustomTestCase): + """Verify correct matching when sequences share suffixes.""" + + def test_shared_suffix_both_match(self): + corpus = _make_corpus("BFS") + corpus.batch_put([[100, 200, 7, 8, 9, 50, 51]]) + corpus.batch_put([[300, 400, 7, 8, 9, 60, 61]]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[7, 8, 9]]) + ids_list = ids.tolist() + self.assertIn(50, ids_list, "Continuation from first sequence missing") + self.assertIn(60, ids_list, "Continuation from second sequence missing") + + +class TestSingleTokenContext(CustomTestCase): + """Verify behavior with minimum-length context.""" + + def test_single_token_query(self): + corpus = _make_corpus("BFS", min_match_window_size=1) + corpus.batch_put([[5, 10, 20, 30]]) + corpus.synchronize() + + ids, masks = corpus.batch_get([[5]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 5, "First token should be last context token") + self.assertIn(10, ids_list, "Should match continuation after single token 5") + + +class TestLongContext(CustomTestCase): + """Verify behavior when query context exceeds branch_length.""" + + def test_context_longer_than_branch_length(self): + corpus = _make_corpus("BFS", branch_length=6, max_match_window_size=5) + seq = list(range(1, 20)) + corpus.batch_put([seq]) + corpus.synchronize() + + long_query = list(range(1, 16)) + ids, masks = corpus.batch_get([long_query]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 15, "First token should be last context token") + self.assertIn(16, ids_list, "Should match via suffix despite long context") + + +class TestDraftBudgetSaturation(CustomTestCase): + """Verify the draft tree uses exactly draft_token_num slots.""" + + def test_full_budget_used(self): + corpus = _make_corpus("BFS", draft_token_num=8) + seq = list(range(1, 30)) + corpus.batch_put([seq]) + corpus.synchronize() + + ids, _ = corpus.batch_get([[1, 2, 3]]) + ids_list = ids.tolist() + self.assertEqual(len(ids_list), 8) + non_zero = [t for t in ids_list[1:] if t != 0] + self.assertGreater( + len(non_zero), + 0, + "Draft budget should have non-zero tokens when cache has long chains", + ) + + +class TestTruncate(CustomTestCase): + """Verify the Result.truncate method via the Python binding.""" + + def test_truncate_reduces_output(self): + corpus = _make_corpus("BFS", draft_token_num=8) + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + + result = corpus._ngram.batchMatch([[1, 2, 3]]) + original_len = len(result.token) + self.assertEqual(original_len, 8) + + result.truncate(4) + self.assertEqual(len(result.token), 4) + self.assertEqual(len(result.mask), 4 * 4) + + def test_truncate_preserves_mask_structure(self): + corpus = _make_corpus("BFS", draft_token_num=8) + corpus.batch_put(SEED_SEQUENCES) + corpus.synchronize() + + result = corpus._ngram.batchMatch([[1, 2, 3]]) + full_ids = list(result.token) + full_mask = list(result.mask) + n = len(full_ids) + + result_copy = corpus._ngram.batchMatch([[1, 2, 3]]) + trunc_n = 4 + result_copy.truncate(trunc_n) + trunc_mask = list(result_copy.mask) + + for i in range(trunc_n): + for j in range(trunc_n): + self.assertEqual( + trunc_mask[i * trunc_n + j], + full_mask[i * n + j], + f"Mask mismatch at ({i},{j})", + ) + + +class TestResetAndReinsert(CustomTestCase): + """Verify that reset followed by new inserts works correctly.""" + + def test_reset_then_reinsert(self): + corpus = _make_corpus("BFS") + corpus.batch_put([[1, 2, 3, 4, 5]]) + corpus.synchronize() + + corpus.reset() + + corpus.batch_put([[10, 20, 30, 40, 50]]) + corpus.synchronize() + + ids_old, _ = corpus.batch_get([[1, 2, 3]]) + ids_old_list = ids_old.tolist() + self.assertTrue( + all(t == 0 for t in ids_old_list[1:]), + f"Old data should not match after reset+reinsert, got {ids_old_list}", + ) + + ids_new, _ = corpus.batch_get([[10, 20, 30]]) + ids_new_list = ids_new.tolist() + self.assertEqual(ids_new_list[0], 30) + self.assertIn(40, ids_new_list, "New data should match after reset+reinsert") + + +class TestSqueezeEvictsOld(CustomTestCase): + """Verify that squeeze actually evicts old data, not just preserves recent.""" + + def test_old_data_evicted(self): + corpus = _make_corpus( + "BFS", capacity=150, branch_length=6, max_match_window_size=5 + ) + + old_seq = list(range(5000, 5030)) + corpus.batch_put([old_seq]) + corpus.synchronize() + + ids_before, _ = corpus.batch_get([[5000, 5001, 5002]]) + self.assertIn( + 5003, + ids_before.tolist(), + "Old data should match before eviction", + ) + + for i in range(5): + new_seq = list(range(6000 + i * 30, 6000 + i * 30 + 30)) + corpus.batch_put([new_seq]) + corpus.synchronize() + + ids_after, _ = corpus.batch_get([[5000, 5001, 5002]]) + ids_after_list = ids_after.tolist() + self.assertNotIn( + 5003, + ids_after_list, + f"Old data should be evicted after pressure, got {ids_after_list}", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3)