diff --git a/docs/advanced_features/speculative_decoding.md b/docs/advanced_features/speculative_decoding.md index 9806f244e..b8fe2d890 100644 --- a/docs/advanced_features/speculative_decoding.md +++ b/docs/advanced_features/speculative_decoding.md @@ -387,7 +387,7 @@ Enable it with: | Parameter | Description | Default | |---|---|---| -| `--speculative-num-draft-tokens` | Number of draft tokens verified per step. | `12` | +| `--speculative-num-draft-tokens` | Number of draft tokens verified per step. If omitted, defaults to `min(--speculative-ngram-max-trie-depth, 12)`. | `12` (with default ngram settings) | | `--speculative-ngram-min-bfs-breadth` | Minimum BFS breadth. | `1` | | `--speculative-ngram-max-bfs-breadth` | Maximum BFS breadth. | `10` | | `--speculative-ngram-match-type` | Ngram tree-building mode: `"BFS"` for recency-based expansion or `"PROB"` for frequency-based expansion. | `"BFS"` | diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp index c79d7cd7d..192b48f6a 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.cpp @@ -63,6 +63,29 @@ void Ngram::asyncInsert(std::vector>&& tokens) { } } +void Ngram::startExternalCorpusLoad() { + std::unique_lock lock(mutex_); + sam_ = std::make_unique(); +} + +void Ngram::appendExternalCorpusTokens(const std::vector& tokens) { + std::unique_lock lock(mutex_); + sam_->appendTokens(tokens); +} + +void Ngram::finishExternalCorpusLoad() { + std::unique_lock lock(mutex_); + sam_->finalize(); + if (sam_->empty()) { + sam_.reset(); + } +} + +void Ngram::clearExternalCorpus() { + std::unique_lock lock(mutex_); + sam_.reset(); +} + void Ngram::insertWorker() { for (;;) { std::vector data; @@ -116,12 +139,17 @@ Result Ngram::batchMatch( std::unique_lock lock(mutex_); - using BuildFn = Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const; - BuildFn build_fn; + using TrieResultBuildFn = + Result (Trie::*)(const int32_t*, size_t, int32_t, size_t, const Param&, MatchState&, size_t) const; + using SamResultBuildFn = Result (SuffixAutomaton::*)(const int32_t*, size_t, int32_t, size_t, const Param&) const; + TrieResultBuildFn trie_result_build_fn; + SamResultBuildFn sam_result_build_fn; if (param_.match_type == "BFS") { - build_fn = &Trie::buildRecency; + trie_result_build_fn = &Trie::buildRecency; + sam_result_build_fn = &SuffixAutomaton::buildRecency; } else if (param_.match_type == "PROB") { - build_fn = &Trie::buildFrequency; + trie_result_build_fn = &Trie::buildFrequency; + sam_result_build_fn = &SuffixAutomaton::buildFrequency; } else { throw std::runtime_error("Unknown match_type: '" + param_.match_type + "'. Must be 'BFS' or 'PROB'."); } @@ -134,9 +162,23 @@ Result Ngram::batchMatch( } auto& state = match_state_[state_ids[i]]; - 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_, state, total_lens[i]); + const auto total_draft_token_num = param_.get_draft_token_num(tokens.size()); + const auto sam_budget = + sam_ && !sam_->empty() ? std::min(param_.external_sam_budget, total_draft_token_num) : size_t{0}; + const auto trie_budget = total_draft_token_num - sam_budget; + + if (sam_budget == 0) { + auto res = (trie_.get()->*trie_result_build_fn)( + suffix.data(), suffix.size(), suffix.back(), total_draft_token_num, param_, state, total_lens[i]); + merged.token.insert(merged.token.end(), res.token.begin(), res.token.end()); + merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end()); + continue; + } + + auto trie_res = (trie_.get()->*trie_result_build_fn)( + suffix.data(), suffix.size(), suffix.back(), trie_budget, param_, state, total_lens[i]); + auto sam_res = (sam_.get()->*sam_result_build_fn)(suffix.data(), suffix.size(), suffix.back(), sam_budget, param_); + auto res = combineRootResults_(suffix.back(), static_cast(total_draft_token_num + 1), trie_res, sam_res); merged.token.insert(merged.token.end(), res.token.begin(), res.token.end()); merged.mask.insert(merged.mask.end(), res.mask.begin(), res.mask.end()); } diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h index d1b404974..fa94d15af 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram.h @@ -3,6 +3,7 @@ #include "param.h" #include "queue.h" #include "result.h" +#include "suffix_automaton.h" #include "trie.h" #include #include @@ -18,6 +19,7 @@ namespace ngram { class Ngram { std::unique_ptr trie_; + std::unique_ptr sam_; Param param_; // NOTE: protects trie_ and pending_count_. Ensures batchMatch never reads @@ -40,6 +42,14 @@ class Ngram { void asyncInsert(std::vector>&& tokens); + void startExternalCorpusLoad(); + + void appendExternalCorpusTokens(const std::vector& tokens); + + void finishExternalCorpusLoad(); + + void clearExternalCorpus(); + Result batchMatch(const std::vector>& tokens); Result batchMatch( diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp index bbd815c51..02f4c3916 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/ngram_corpus_ffi.cpp @@ -23,7 +23,9 @@ struct NgramCorpusObj : public tvm::ffi::Object { int64_t min_bfs_breadth, int64_t max_bfs_breadth, int64_t draft_token_num, - int64_t match_type) { + int64_t match_type, + int64_t external_sam_budget, + int64_t external_corpus_max_tokens) { ngram::Param param; param.enable = true; param.enable_router_mode = false; @@ -32,6 +34,8 @@ struct NgramCorpusObj : public tvm::ffi::Object { param.max_bfs_breadth = static_cast(max_bfs_breadth); param.draft_token_num = static_cast(draft_token_num); param.match_type = (match_type == 0) ? "BFS" : "PROB"; + param.external_sam_budget = static_cast(external_sam_budget); + param.external_corpus_max_tokens = static_cast(external_corpus_max_tokens); ngram_ = std::make_unique(static_cast(capacity), param); } @@ -97,6 +101,25 @@ struct NgramCorpusObj : public tvm::ffi::Object { ngram_->eraseMatchState(state_ids); } + void start_external_corpus_load() { + ngram_->startExternalCorpusLoad(); + } + + void append_external_corpus_tokens(const tvm::ffi::TensorView tokens_tv) { + auto* data = static_cast(tokens_tv.data_ptr()); + int64_t n = tokens_tv.size(0); + std::vector tokens(data, data + n); + ngram_->appendExternalCorpusTokens(tokens); + } + + void finish_external_corpus_load() { + ngram_->finishExternalCorpusLoad(); + } + + void clear_external_corpus() { + ngram_->clearExternalCorpus(); + } + void synchronize() { ngram_->synchronize(); } @@ -130,11 +153,15 @@ struct NgramCorpusObj : public tvm::ffi::Object { void register_ngram_corpus() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() - .def(refl::init(), "__init__") + .def(refl::init(), "__init__") .def("async_insert", &NgramCorpusObj::async_insert) .def("batch_match", &NgramCorpusObj::batch_match) .def("batch_match_stateful", &NgramCorpusObj::batch_match_stateful) .def("erase_match_state", &NgramCorpusObj::erase_match_state) + .def("start_external_corpus_load", &NgramCorpusObj::start_external_corpus_load) + .def("append_external_corpus_tokens", &NgramCorpusObj::append_external_corpus_tokens) + .def("finish_external_corpus_load", &NgramCorpusObj::finish_external_corpus_load) + .def("clear_external_corpus", &NgramCorpusObj::clear_external_corpus) .def("synchronize", &NgramCorpusObj::synchronize) .def("reset", &NgramCorpusObj::reset); } diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/param.h b/python/sglang/jit_kernel/csrc/ngram_corpus/param.h index 725f635db..9c2701b1b 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/param.h +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/param.h @@ -19,6 +19,8 @@ struct Param { size_t max_bfs_breadth; size_t max_trie_depth; size_t draft_token_num; + size_t external_sam_budget = 0; + size_t external_corpus_max_tokens = 10000000; std::string match_type; std::vector batch_draft_token_num; @@ -92,7 +94,8 @@ struct Param { ss << "enable = " << enable << ", enable_router_mode = " << enable_router_mode << ", min_bfs_breadth = " << min_bfs_breadth << ", max_bfs_breadth = " << max_bfs_breadth << ", max_trie_depth = " << max_trie_depth << ", draft_token_num = " << draft_token_num - << ", match_type = " << match_type; + << ", external_sam_budget = " << external_sam_budget + << ", external_corpus_max_tokens = " << external_corpus_max_tokens << ", match_type = " << match_type; ss << ", batch_draft_token_num(" << batch_draft_token_num.size() << ") = "; for (int i = 0; i < batch_draft_token_num.size(); ++i) { ss << i << "|" << batch_draft_token_num[i] << ","; diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp index 404b7a3f2..07138bf8d 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/result.cpp @@ -1,5 +1,6 @@ #include "result.h" +#include #include #include #include @@ -47,6 +48,81 @@ Result fillResult(int last_token, int draft_token_num, std::vector& tree, return info; } +std::vector> extractLeafPaths_(const Result& result) { + const auto n = static_cast(result.token.size()); + if (n <= 1) { + return {}; + } + + std::vector parent(n, -1); + std::vector has_child(n, false); + for (int i = 1; i < n; ++i) { + for (int j = i - 1; j >= 0; --j) { + if (result.mask[i * n + j]) { + parent[i] = j; + has_child[j] = true; + break; + } + } + } + + std::vector> paths; + for (int leaf = 1; leaf < n; ++leaf) { + if (has_child[leaf]) { + continue; + } + std::vector path; + for (int cursor = leaf; cursor > 0; cursor = parent[cursor]) { + path.emplace_back(result.token[cursor]); + } + std::reverse(path.begin(), path.end()); + if (path.size() == 1 && path.front() == 0) { + continue; + } + paths.emplace_back(std::move(path)); + } + return paths; +} + +Result buildResultFromLeafPaths_(int last_token, int draft_token_num, const std::vector>& paths) { + std::vector tree(draft_token_num); + const int root = 0; + int cursor = 1; + for (const auto& path : paths) { + int parent = root; + for (const auto token : path) { + auto iter = tree[parent].next.find(token); + if (iter == tree[parent].next.end()) { + if (cursor >= draft_token_num) { + parent = -1; + break; + } + iter = tree[parent].next.insert({token, cursor++}).first; + } + parent = iter->second; + } + if (cursor >= draft_token_num) { + break; + } + } + return fillResult(last_token, draft_token_num, tree, root); +} + +Result combineRootResults_(int last_token, int draft_token_num, const Result& primary, const Result& secondary) { + auto primary_paths = extractLeafPaths_(primary); + auto secondary_paths = extractLeafPaths_(secondary); + std::vector> merged_paths = std::move(primary_paths); + merged_paths.reserve(merged_paths.size() + secondary_paths.size()); + for (const auto& path : secondary_paths) { + if (path.empty()) { + continue; + } + merged_paths.emplace_back(path); + } + + return buildResultFromLeafPaths_(last_token, draft_token_num, merged_paths); +} + void Result::truncate(size_t n) { if (n < token.size()) { int full_n = token.size(); diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/result.h b/python/sglang/jit_kernel/csrc/ngram_corpus/result.h index c48351d77..3e7cc6a82 100644 --- a/python/sglang/jit_kernel/csrc/ngram_corpus/result.h +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/result.h @@ -18,5 +18,8 @@ struct Node { }; Result fillResult(int last_token, int draft_token_num, std::vector& tree, int root); +std::vector> extractLeafPaths_(const Result& result); +Result buildResultFromLeafPaths_(int last_token, int draft_token_num, const std::vector>& paths); +Result combineRootResults_(int last_token, int draft_token_num, const Result& primary, const Result& secondary); } // namespace ngram diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.cpp b/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.cpp new file mode 100644 index 000000000..65d8a7d1b --- /dev/null +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.cpp @@ -0,0 +1,283 @@ +#include "suffix_automaton.h" + +#include +#include +#include +#include +#include +#include + +namespace ngram { + +SuffixAutomaton::SuffixAutomaton() { + reset_(); +} + +void SuffixAutomaton::reset_() { + states_.clear(); + states_.emplace_back(); + last_ = 0; + pos_ = 0; + saw_token_ = false; + finalized_ = false; + loaded_ = false; +} + +void SuffixAutomaton::appendTokens(const std::vector& tokens) { + if (finalized_) { + throw std::runtime_error("Cannot append tokens after finalizing the SAM."); + } + if (tokens.empty()) { + return; + } + + for (const auto token : tokens) { + extend_(token, pos_++); + saw_token_ = true; + } +} + +void SuffixAutomaton::finalize() { + if (finalized_) { + return; + } + finalized_ = true; + if (!saw_token_) { + return; + } + + propagateOccurrencesAndRecency_(); + loaded_ = true; +} + +void SuffixAutomaton::extend_(int32_t token, int64_t pos) { + const int cur = static_cast(states_.size()); + states_.emplace_back(); + states_[cur].max_len = states_[last_].max_len + 1; + states_[cur].occ_count = 1; + states_[cur].max_end_pos = pos; + + int p = last_; + while (p != -1 && !states_[p].next.contains(token)) { + states_[p].next[token] = cur; + p = states_[p].link; + } + + if (p == -1) { + states_[cur].link = 0; + last_ = cur; + return; + } + + const int q = states_[p].next[token]; + if (states_[p].max_len + 1 == states_[q].max_len) { + states_[cur].link = q; + last_ = cur; + return; + } + + const int clone = static_cast(states_.size()); + states_.push_back(states_[q]); + states_[clone].max_len = states_[p].max_len + 1; + states_[clone].occ_count = 0; + states_[clone].children_by_freq.clear(); + states_[clone].children_by_recency.clear(); + + while (p != -1 && states_[p].next[token] == q) { + states_[p].next[token] = clone; + p = states_[p].link; + } + + states_[q].link = clone; + states_[cur].link = clone; + last_ = cur; +} + +void SuffixAutomaton::propagateOccurrencesAndRecency_() { + std::vector order(states_.size()); + std::iota(order.begin(), order.end(), 0); + std::sort( + order.begin(), order.end(), [this](int lhs, int rhs) { return states_[lhs].max_len < states_[rhs].max_len; }); + + for (auto it = order.rbegin(); it != order.rend(); ++it) { + const int state = *it; + const int link = states_[state].link; + if (link < 0) { + continue; + } + states_[link].occ_count += states_[state].occ_count; + states_[link].max_end_pos = std::max(states_[link].max_end_pos, states_[state].max_end_pos); + } + + for (auto& state : states_) { + state.children_by_freq.clear(); + state.children_by_recency.clear(); + state.children_by_freq.reserve(state.next.size()); + state.children_by_recency.reserve(state.next.size()); + for (const auto& [token, child_state] : state.next) { + if (token == kSeparatorToken) { + continue; + } + state.children_by_freq.emplace_back(token, child_state); + state.children_by_recency.emplace_back(token, child_state); + } + + std::sort(state.children_by_freq.begin(), state.children_by_freq.end(), [this](const auto& lhs, const auto& rhs) { + const auto lhs_freq = states_[lhs.second].occ_count; + const auto rhs_freq = states_[rhs.second].occ_count; + return std::tie(rhs_freq, lhs.first, lhs.second) < std::tie(lhs_freq, rhs.first, rhs.second); + }); + std::sort( + state.children_by_recency.begin(), state.children_by_recency.end(), [this](const auto& lhs, const auto& rhs) { + const auto lhs_recency = states_[lhs.second].max_end_pos; + const auto rhs_recency = states_[rhs.second].max_end_pos; + return std::tie(rhs_recency, lhs.first, lhs.second) < std::tie(lhs_recency, rhs.first, rhs.second); + }); + } +} + +std::vector SuffixAutomaton::match(const int32_t* context, size_t len, size_t max_depth) const { + if (empty() || len == 0) { + return {}; + } + + const auto start = len > max_depth ? len - max_depth : 0; + int state = 0; + int32_t matched_len = 0; + for (size_t i = start; i < len; ++i) { + const auto token = context[i]; + while (state != 0 && !states_[state].next.contains(token)) { + state = states_[state].link; + matched_len = std::min(matched_len, states_[state].max_len); + } + if (auto iter = states_[state].next.find(token); iter != states_[state].next.end()) { + state = iter->second; + ++matched_len; + } else if (auto root_iter = states_[0].next.find(token); root_iter != states_[0].next.end()) { + state = root_iter->second; + matched_len = 1; + } else { + state = 0; + matched_len = 0; + } + } + + std::vector anchors; + while (state > 0 && matched_len > 0) { + if (!states_[state].children_by_freq.empty()) { + anchors.push_back({state, matched_len}); + } + state = states_[state].link; + if (state <= 0) { + break; + } + matched_len = std::min(matched_len, states_[state].max_len); + } + return anchors; +} + +Result SuffixAutomaton::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.max_trie_depth); + const auto max_match_depth = std::max(1, static_cast(param.max_trie_depth - 1)); + const double bfs_breadth_scale = double(param.max_bfs_breadth - param.min_bfs_breadth) / max_match_depth; + std::vector tree(draft_token_num + 1); + int root = 0; + int cursor = 1; + + for (const auto& anchor : anchors) { + std::queue> queue; + queue.push( + {root, (max_match_depth - anchor.matched_len) * bfs_breadth_scale + param.min_bfs_breadth, anchor.state}); + while (!queue.empty() && cursor <= static_cast(draft_token_num)) { + auto [parent, cur_breadth, state] = queue.front(); + queue.pop(); + + const auto& children = states_[state].children_by_recency; + const auto breadth = std::max(1, static_cast(cur_breadth)); + for (int i = 0; + i < breadth && i < static_cast(children.size()) && cursor <= static_cast(draft_token_num); + ++i) { + const auto [token, child_state] = children[i]; + int pos = -1; + if (auto iter = tree[parent].next.find(token); iter != tree[parent].next.end()) { + pos = iter->second; + } else { + pos = tree[parent].next.insert({token, cursor++}).first->second; + } + queue.emplace(pos, cur_breadth - bfs_breadth_scale, child_state); + } + } + } + return fillResult(last_token, draft_token_num + 1, tree, root); +} + +Result SuffixAutomaton::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.max_trie_depth); + struct CompareByProb { + bool operator()( + const std::tuple& lhs, const std::tuple& rhs) const { + return std::get<3>(lhs) < std::get<3>(rhs); + } + }; + + std::priority_queue< + std::tuple, + std::vector>, + CompareByProb> + heap; + std::vector tree(draft_token_num + 1); + int root = 0; + int cursor = 1; + const int top_k = static_cast(param.max_bfs_breadth); + + auto addToHeap = [this, &heap, top_k](int parent, int state, double prob) { + if (top_k <= 0) { + return; + } + const auto& children = states_[state].children_by_freq; + if (children.empty()) { + return; + } + double sum_freq = 0.0; + int count = 0; + for (const auto& [_, child_state] : children) { + sum_freq += static_cast(states_[child_state].occ_count); + if (++count >= top_k) { + break; + } + } + if (sum_freq <= 0) { + sum_freq = 1.0; + } + count = 0; + for (const auto& [token, child_state] : children) { + const auto scaled_prob = static_cast(states_[child_state].occ_count) / sum_freq * prob; + heap.emplace(parent, token, child_state, scaled_prob); + if (++count >= top_k) { + break; + } + } + }; + + for (const auto& anchor : anchors) { + addToHeap(root, anchor.state, 1.0); + while (!heap.empty() && cursor <= static_cast(draft_token_num)) { + auto [parent, token, child_state, prob] = heap.top(); + heap.pop(); + + int pos = -1; + if (auto iter = tree[parent].next.find(token); iter != tree[parent].next.end()) { + pos = iter->second; + } else { + pos = cursor++; + tree[parent].next[token] = pos; + } + addToHeap(pos, child_state, prob); + } + } + return fillResult(last_token, draft_token_num + 1, tree, root); +} + +} // namespace ngram diff --git a/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.h b/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.h new file mode 100644 index 000000000..ebd8f5471 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/ngram_corpus/suffix_automaton.h @@ -0,0 +1,62 @@ +#pragma once + +#include "param.h" +#include "result.h" +#include +#include +#include +#include +#include + +namespace ngram { + +struct SamAnchor { + int state = 0; + int32_t matched_len = 0; +}; + +struct SamState { + int link = -1; + int32_t max_len = 0; + std::unordered_map next; + uint64_t occ_count = 0; + int64_t max_end_pos = -1; + std::vector> children_by_freq; + std::vector> children_by_recency; +}; + +class SuffixAutomaton { + public: + static constexpr int32_t kSeparatorToken = std::numeric_limits::min(); + + SuffixAutomaton(); + + void appendTokens(const std::vector& tokens); + + void finalize(); + + bool empty() const { + return !loaded_; + } + + 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; + + private: + void reset_(); + void extend_(int32_t token, int64_t pos); + void propagateOccurrencesAndRecency_(); + std::vector match(const int32_t* context, size_t len, size_t max_depth) const; + + std::vector states_; + int last_ = 0; + int64_t pos_ = 0; + bool saw_token_ = false; + bool finalized_ = false; + bool loaded_ = false; +}; + +} // namespace ngram diff --git a/python/sglang/jit_kernel/ngram_corpus.py b/python/sglang/jit_kernel/ngram_corpus.py index 2182c5c34..2a6ba13e8 100644 --- a/python/sglang/jit_kernel/ngram_corpus.py +++ b/python/sglang/jit_kernel/ngram_corpus.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Iterable, Sequence from typing import List, Tuple import numpy as np @@ -29,6 +30,7 @@ def get_ngram_corpus_cls(): cpp_files=[ "ngram_corpus/result.cpp", "ngram_corpus/trie.cpp", + "ngram_corpus/suffix_automaton.cpp", "ngram_corpus/ngram.cpp", "ngram_corpus/ngram_corpus_ffi.cpp", ], @@ -48,6 +50,8 @@ def get_ngram_corpus_cls(): max_bfs_breadth: int, draft_token_num: int, match_type: str, + external_sam_budget: int = 0, + external_corpus_max_tokens: int = 10000000, ) -> None: mt = _MATCH_TYPE_MAP.get(match_type) if mt is None: @@ -61,6 +65,8 @@ def get_ngram_corpus_cls(): max_bfs_breadth, draft_token_num, mt, + external_sam_budget, + external_corpus_max_tokens, ) self._draft_token_num = draft_token_num @@ -112,4 +118,22 @@ def get_ngram_corpus_cls(): state_ids_t = torch.tensor(state_ids, dtype=torch.int64) self.erase_match_state(state_ids_t) # type: ignore + def load_external_corpus( + self, chunks: Iterable[Sequence[int]] + ) -> Tuple[int, int]: + self.start_external_corpus_load() # type: ignore + chunk_count = 0 + loaded_token_count = 0 + try: + for chunk in chunks: + tokens_t = torch.tensor(list(chunk), dtype=torch.int32) + loaded_token_count += len(tokens_t) + self.append_external_corpus_tokens(tokens_t) # type: ignore + chunk_count += 1 + self.finish_external_corpus_load() # type: ignore + except Exception: + self.clear_external_corpus() # type: ignore + raise + return chunk_count, loaded_token_count + return NgramCorpusFFI diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d91ced805..646edb414 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -512,6 +512,9 @@ class ServerArgs: speculative_ngram_match_type: Literal["BFS", "PROB"] = "BFS" speculative_ngram_max_trie_depth: int = 18 speculative_ngram_capacity: int = 10 * 1000 * 1000 + speculative_ngram_external_corpus_path: Optional[str] = None + speculative_ngram_external_sam_budget: int = 0 + speculative_ngram_external_corpus_max_tokens: int = 10000000 enable_multi_layer_eagle: bool = False # Expert parallelism @@ -3110,6 +3113,30 @@ class ServerArgs: "speculative_num_draft_tokens is set to 12 by default for ngram speculative decoding. " "You can override this by explicitly setting --speculative-num-draft-tokens." ) + if self.speculative_ngram_external_corpus_path is not None: + if self.speculative_ngram_external_sam_budget <= 0: + raise ValueError( + "--speculative-ngram-external-sam-budget must be positive when " + "--speculative-ngram-external-corpus-path is set." + ) + if self.speculative_ngram_external_corpus_max_tokens <= 0: + raise ValueError( + "--speculative-ngram-external-corpus-max-tokens must be positive when " + "--speculative-ngram-external-corpus-path is set." + ) + if ( + self.speculative_ngram_external_sam_budget + > self.speculative_num_draft_tokens - 1 + ): + raise ValueError( + "speculative_ngram_external_sam_budget must be less than or equal to " + f"speculative_num_draft_tokens - 1 ({self.speculative_num_draft_tokens - 1})." + ) + elif self.speculative_ngram_external_sam_budget != 0: + raise ValueError( + "--speculative-ngram-external-sam-budget requires " + "--speculative-ngram-external-corpus-path." + ) logger.warning( "The overlap scheduler and mixed chunked prefill are disabled because of " "using ngram speculative decoding." @@ -4888,6 +4915,24 @@ class ServerArgs: default=ServerArgs.speculative_ngram_capacity, help="The cache capacity for ngram speculative decoding.", ) + parser.add_argument( + "--speculative-ngram-external-corpus-path", + type=str, + default=ServerArgs.speculative_ngram_external_corpus_path, + help="Optional path to an external corpus used to build a read-only SAM for ngram speculative decoding.", + ) + parser.add_argument( + "--speculative-ngram-external-sam-budget", + type=int, + default=ServerArgs.speculative_ngram_external_sam_budget, + help="Number of draft nodes reserved for the external SAM subtree in ngram speculative decoding.", + ) + parser.add_argument( + "--speculative-ngram-external-corpus-max-tokens", + type=int, + default=ServerArgs.speculative_ngram_external_corpus_max_tokens, + help="Fail startup if the tokenized external ngram corpus exceeds this many tokens. Tune this based on your CPU memory budget.", + ) # Multi-layer Eagle speculative decoding parser.add_argument( diff --git a/python/sglang/srt/speculative/cpp_ngram/external_corpus.py b/python/sglang/srt/speculative/cpp_ngram/external_corpus.py new file mode 100644 index 000000000..62445adde --- /dev/null +++ b/python/sglang/srt/speculative/cpp_ngram/external_corpus.py @@ -0,0 +1,62 @@ +import json +from collections.abc import Iterator +from pathlib import Path + +# Must match SuffixAutomaton::kSeparatorToken in suffix_automaton.h. +SEPARATOR_TOKEN = -(2**31) + +# Default chunk size for streaming tokenized documents into the SAM. +DEFAULT_CHUNK_SIZE = 4096 + + +def iter_external_corpus_chunks( + path: str, tokenizer, max_tokens: int, chunk_size: int = DEFAULT_CHUNK_SIZE +) -> Iterator[list[int]]: + """Chunk documents and yield fixed-size token chunks from a JSONL corpus file.""" + corpus_path = Path(path) + if not corpus_path.is_file(): + raise ValueError(f"External ngram corpus path does not exist: {path}") + if tokenizer is None: + raise ValueError("A tokenizer is required to load an external ngram corpus.") + if max_tokens <= 0: + raise ValueError("External ngram corpus max tokens must be positive.") + + total_tokens = 0 + has_previous_doc = False + with corpus_path.open("r", encoding="utf-8") as f: + for line_no, line in enumerate(f, start=1): + if not line.strip(): + continue + + try: + record = json.loads(line) + except json.JSONDecodeError as e: + raise ValueError( + f"Invalid JSON in external ngram corpus at line {line_no}: {e.msg}" + ) from e + + if not isinstance(record, str): + raise ValueError( + "Invalid external ngram corpus record at line " + f"{line_no}: expected a JSON string." + ) + + token_ids = list(tokenizer.encode(record, add_special_tokens=False)) + if not token_ids: + continue + + separator_cost = 1 if has_previous_doc else 0 + next_total_tokens = total_tokens + separator_cost + len(token_ids) + if next_total_tokens > max_tokens: + raise ValueError( + "External ngram corpus exceeds the configured token limit " + f"({max_tokens}) at line {line_no} after loading " + f"{total_tokens} tokens." + ) + total_tokens = next_total_tokens + + if has_previous_doc: + token_ids = [SEPARATOR_TOKEN] + token_ids + for i in range(0, len(token_ids), chunk_size): + yield token_ids[i : i + chunk_size] + has_previous_doc = True diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py index 0795e3bf4..b7d5f2a9b 100644 --- a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py +++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py @@ -1,15 +1,44 @@ # -*- coding: utf-8 -*- import logging -from typing import Dict, List, Tuple +from collections.abc import Iterable, Iterator, Sequence +from typing import Dict, List, Optional, Tuple import numpy as np from sglang.jit_kernel.ngram_corpus import get_ngram_corpus_cls +from sglang.srt.speculative.cpp_ngram.external_corpus import SEPARATOR_TOKEN logger = logging.getLogger(__name__) +# Convenience path for pre-tokenized in-memory documents. The main serving +# startup path loads file-backed corpora through `iter_external_corpus_chunks()`. +def _documents_to_chunks( + documents: Iterable[Sequence[int]], + max_tokens: Optional[int] = None, +) -> Iterator[list[int]]: + total_tokens = 0 + has_previous = False + for doc in documents: + if not doc: + continue + doc_tokens = list(doc) + separator_cost = 1 if has_previous else 0 + next_total_tokens = total_tokens + separator_cost + len(doc_tokens) + if max_tokens is not None and next_total_tokens > max_tokens: + raise ValueError( + "External ngram corpus exceeds the configured token limit " + f"({max_tokens}) after loading {total_tokens} tokens." + ) + total_tokens = next_total_tokens + if has_previous: + yield [SEPARATOR_TOKEN] + doc_tokens + else: + yield doc_tokens + has_previous = True + + class NgramCorpus: def __init__( self, @@ -19,6 +48,9 @@ class NgramCorpus: draft_token_num=8, match_type="BFS", capacity=1000000, + external_sam_budget=0, + external_corpus_max_tokens=10000000, + external_corpus_documents: Optional[Iterable[Sequence[int]]] = None, ) -> None: cls = get_ngram_corpus_cls() self._obj = cls( @@ -28,11 +60,20 @@ class NgramCorpus: max_bfs_breadth=max_bfs_breadth, draft_token_num=draft_token_num, match_type=match_type, + external_sam_budget=external_sam_budget, + external_corpus_max_tokens=external_corpus_max_tokens, ) self.default_mask = np.ones((1, 1), dtype=np.int64) self.draft_token_num = draft_token_num self._req_id_to_state_id: Dict[str, int] = {} self._next_state_id: int = 0 + self.external_corpus_token_count = 0 + if external_corpus_documents is not None: + self.load_external_corpus( + _documents_to_chunks( + external_corpus_documents, external_corpus_max_tokens + ) + ) def _get_state_id(self, req_id: str) -> int: sid = self._req_id_to_state_id.get(req_id) @@ -48,6 +89,18 @@ class NgramCorpus: def synchronize(self): self._obj.synchronize() # type: ignore + def load_external_corpus(self, chunks: Iterable[Sequence[int]]) -> int: + """Load pre-chunked external corpus tokens. + + Callers passing raw chunk iterables must enforce any token budget before + calling this method. Python-side helpers such as + `iter_external_corpus_chunks()` and `external_corpus_documents=` handle + `external_corpus_max_tokens` validation before handing chunks to C++. + """ + _, loaded_token_count = self._obj.load_external_corpus(chunks) + self.external_corpus_token_count = loaded_token_count + return loaded_token_count + def reset(self): self._obj.reset() # type: ignore self._req_id_to_state_id.clear() diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index cbc958660..afc761a2d 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -11,6 +11,9 @@ 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.external_corpus import ( + iter_external_corpus_chunks, +) 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 @@ -54,7 +57,21 @@ class NGRAMWorker: capacity=server_args.speculative_ngram_capacity, max_trie_depth=server_args.speculative_ngram_max_trie_depth, draft_token_num=server_args.speculative_num_draft_tokens, + external_sam_budget=server_args.speculative_ngram_external_sam_budget, + external_corpus_max_tokens=server_args.speculative_ngram_external_corpus_max_tokens, ) + if server_args.speculative_ngram_external_corpus_path is not None: + loaded_token_count = self.ngram_corpus.load_external_corpus( + iter_external_corpus_chunks( + server_args.speculative_ngram_external_corpus_path, + target_worker.tokenizer, + server_args.speculative_ngram_external_corpus_max_tokens, + ) + ) + logger.info( + "Loaded external ngram corpus (%d tokens) for SAM speculative decoding.", + loaded_token_count, + ) def clear_cache_pool(self): self.ngram_corpus.reset() diff --git a/test/registered/spec/test_ngram_speculative_decoding.py b/test/registered/spec/test_ngram_speculative_decoding.py index 63c75b116..690e6ec45 100644 --- a/test/registered/spec/test_ngram_speculative_decoding.py +++ b/test/registered/spec/test_ngram_speculative_decoding.py @@ -1,5 +1,10 @@ +import json +import os +import tempfile import unittest +import requests + from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci @@ -30,6 +35,21 @@ DEFAULT_SERVER_ARGS = [ 0.8, ] +EXTERNAL_SAM_CORPUS_RECORDS = [ + "The capital of France is Paris.", + "The answer to life, the universe, and everything is 42.", +] + + +def _safe_remove(path: str): + if os.path.exists(path): + os.remove(path) + + +def _safe_kill_process(process): + if process is not None and process.poll() is None: + kill_process_tree(process.pid) + class TestNgramSpeculativeDecodingBase(GSM8KMixin, CustomTestCase): model = DEFAULT_TARGET_MODEL_NGRAM @@ -86,5 +106,66 @@ class TestNgramSpeculativeDecodingPaged(TestNgramSpeculativeDecodingBase): ] +class TestNgramExternalSamSmoke(CustomTestCase): + model = DEFAULT_TARGET_MODEL_NGRAM + base_url = DEFAULT_URL_FOR_TEST + attention_backends = ("triton", "flashinfer") + + def get_server_args(self, attention_backend): + return DEFAULT_SERVER_ARGS + [ + "--attention-backend", + attention_backend, + "--speculative-ngram-external-corpus-path", + self.external_corpus_path, + "--speculative-ngram-external-sam-budget", + "4", + ] + + @classmethod + def setUpClass(cls): + envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.set(False) + envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False) + with tempfile.NamedTemporaryFile( + mode="w", suffix=".jsonl", prefix="ngram_external_sam_", delete=False + ) as f: + for record in EXTERNAL_SAM_CORPUS_RECORDS: + f.write(json.dumps(record)) + f.write("\n") + cls.external_corpus_path = f.name + cls.addClassCleanup(_safe_remove, cls.external_corpus_path) + + def _run_external_sam_smoke(self, attention_backend): + process = popen_launch_server( + self.model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=self.get_server_args(attention_backend), + ) + try: + response = requests.post( + self.base_url + "/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 8, + }, + }, + timeout=120, + ) + self.assertEqual(response.status_code, 200, response.text) + response_json = response.json() + self.assertIn("text", response_json) + self.assertIn("meta_info", response_json) + self.assertGreater(response_json["meta_info"]["completion_tokens"], 0) + finally: + _safe_kill_process(process) + + def test_generate_with_external_sam(self): + for attention_backend in self.attention_backends: + with self.subTest(attention_backend=attention_backend): + self._run_external_sam_smoke(attention_backend) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index ff381bd21..d18c5ab99 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -426,5 +426,63 @@ class TestHiCacheArgs(unittest.TestCase): self.assertEqual(args.decode_attention_backend, "triton") +class TestNgramExternalSamArgs(CustomTestCase): + def _make_dummy_ngram_args(self, **overrides): + args = ServerArgs(model_path="dummy") + args.speculative_algorithm = "NGRAM" + args.speculative_num_draft_tokens = 12 + args.device = "cuda" + for key, value in overrides.items(): + setattr(args, key, value) + return args + + def test_prepare_server_args_parses_external_sam_args(self): + server_args = prepare_server_args( + [ + "--model-path", + "dummy", + "--speculative-algorithm", + "NGRAM", + "--speculative-ngram-external-corpus-path", + "/tmp/ngram-corpus.jsonl", + "--speculative-ngram-external-sam-budget", + "4", + "--speculative-ngram-external-corpus-max-tokens", + "128", + ] + ) + self.assertEqual( + server_args.speculative_ngram_external_corpus_path, + "/tmp/ngram-corpus.jsonl", + ) + self.assertEqual(server_args.speculative_ngram_external_sam_budget, 4) + self.assertEqual(server_args.speculative_ngram_external_corpus_max_tokens, 128) + + def test_external_sam_budget_requires_path(self): + with self.assertRaises(ValueError) as context: + self._make_dummy_ngram_args( + speculative_ngram_external_sam_budget=2, + )._handle_speculative_decoding() + self.assertIn("external-sam-budget", str(context.exception)) + + def test_external_sam_budget_must_fit_draft_budget(self): + with self.assertRaises(ValueError) as context: + self._make_dummy_ngram_args( + speculative_num_draft_tokens=4, + speculative_ngram_external_corpus_path="/tmp/ngram-corpus.jsonl", + speculative_ngram_external_sam_budget=4, + )._handle_speculative_decoding() + self.assertIn("speculative_num_draft_tokens - 1", str(context.exception)) + + def test_external_corpus_max_tokens_must_be_positive(self): + with self.assertRaises(ValueError) as context: + self._make_dummy_ngram_args( + speculative_ngram_external_corpus_path="/tmp/ngram-corpus.jsonl", + speculative_ngram_external_sam_budget=2, + speculative_ngram_external_corpus_max_tokens=0, + )._handle_speculative_decoding() + self.assertIn("external-corpus-max-tokens", str(context.exception)) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/spec/test_ngram_corpus.py b/test/registered/unit/spec/test_ngram_corpus.py index 94a157183..2e652f260 100644 --- a/test/registered/unit/spec/test_ngram_corpus.py +++ b/test/registered/unit/spec/test_ngram_corpus.py @@ -1,8 +1,14 @@ +import json +import os +import tempfile import unittest import uuid import numpy as np +from sglang.srt.speculative.cpp_ngram.external_corpus import ( + iter_external_corpus_chunks, +) from sglang.srt.speculative.cpp_ngram.ngram_corpus import NgramCorpus from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -17,6 +23,9 @@ def _make_corpus(match_type="BFS", **kwargs): max_bfs_breadth=8, draft_token_num=8, capacity=100000, + external_sam_budget=0, + external_corpus_max_tokens=10000000, + external_corpus_documents=None, ) defaults.update(kwargs) defaults["match_type"] = match_type @@ -43,6 +52,12 @@ def _batch_get_with_state( return corpus.batch_get([req_id], [current_tokens], [total_len]) +class _IntTokenizer: + def encode(self, text: str, add_special_tokens: bool = False): + del add_special_tokens + return [int(piece) for piece in text.split()] + + SEED_SEQUENCES = [ [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 44, 55, 66, 77, 88, 99, 100], @@ -674,6 +689,159 @@ class TestNgramCorpusIncremental(CustomTestCase): np.testing.assert_array_equal(inc_masks, full_masks) +class TestNgramCorpusExternalSam(CustomTestCase): + """Verify external SAM loading and fixed-budget composition.""" + + def test_external_corpus_iterator_streams_documents(self): + corpus = _make_corpus( + "BFS", + draft_token_num=4, + external_sam_budget=3, + external_corpus_max_tokens=8, + ) + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + f.write(json.dumps("1 2 3 4 5")) + f.write("\n") + f.write(json.dumps("8 9")) + f.write("\n") + path = f.name + self.addCleanup(os.remove, path) + + loaded_token_count = corpus.load_external_corpus( + iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=8) + ) + # 5 doc tokens + 1 separator + 2 doc tokens = 8 + self.assertEqual(loaded_token_count, 8) + + ids, _ = _batch_get(corpus, [[1, 2, 3]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 3) + self.assertEqual(ids_list[1:3], [4, 5]) + + def test_external_corpus_iterator_rejects_oversized_corpus(self): + corpus = _make_corpus( + "BFS", + external_sam_budget=2, + external_corpus_max_tokens=4, + ) + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + f.write(json.dumps("1 2 3")) + f.write("\n") + f.write(json.dumps("4 5")) + f.write("\n") + path = f.name + self.addCleanup(os.remove, path) + + with self.assertRaisesRegex(ValueError, "token limit"): + corpus.load_external_corpus( + iter_external_corpus_chunks(path, _IntTokenizer(), max_tokens=4) + ) + + def test_external_sam_documents_reject_oversized_corpus(self): + with self.assertRaisesRegex(ValueError, "token limit"): + _make_corpus( + "BFS", + external_sam_budget=2, + external_corpus_max_tokens=4, + external_corpus_documents=[[1, 2, 3], [4, 5]], + ) + + def test_external_sam_only_chain(self): + corpus = _make_corpus( + "BFS", + draft_token_num=4, + external_sam_budget=3, + external_corpus_documents=[[1, 2, 3, 4, 5]], + ) + + ids, masks = _batch_get(corpus, [[1, 2, 3]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 3) + self.assertEqual(ids_list[1:3], [4, 5]) + + def test_external_sam_respects_document_boundaries(self): + corpus = _make_corpus( + "BFS", + draft_token_num=4, + external_sam_budget=3, + external_corpus_documents=[[1, 2, 3], [4, 5, 6]], + ) + + ids, _ = _batch_get(corpus, [[2, 3]]) + ids_list = ids.tolist() + self.assertEqual(ids_list[0], 3) + self.assertTrue(all(token == 0 for token in ids_list[1:]), ids_list) + + def test_external_sam_adds_distinct_root_branch(self): + corpus = _make_corpus( + "BFS", + draft_token_num=6, + external_sam_budget=2, + external_corpus_documents=[[1, 2, 3, 20, 21]], + ) + corpus.batch_put([[1, 2, 3, 10, 11]]) + corpus.synchronize() + + ids, masks = _batch_get(corpus, [[1, 2, 3]]) + leaf_paths = corpus.leaf_paths_from_mask( + ids.tolist(), masks.reshape(6, 6).tolist() + ) + self.assertIn([3, 10, 11], leaf_paths) + self.assertIn([3, 20, 21], leaf_paths) + + def test_shared_prefix_keeps_both_branches(self): + corpus = _make_corpus( + "BFS", + draft_token_num=5, + external_sam_budget=2, + external_corpus_documents=[[1, 2, 3, 10, 99]], + ) + corpus.batch_put([[1, 2, 3, 10, 11]]) + corpus.synchronize() + + ids, masks = _batch_get(corpus, [[1, 2, 3]]) + leaf_paths = corpus.leaf_paths_from_mask( + ids.tolist(), masks.reshape(5, 5).tolist() + ) + self.assertIn([3, 10, 11], leaf_paths) + self.assertIn([3, 10, 99], leaf_paths) + + def test_shared_prefix_merge_can_underfill_budget(self): + corpus = _make_corpus( + "BFS", + draft_token_num=6, + external_sam_budget=2, + external_corpus_documents=[[1, 2, 3, 10, 99]], + ) + corpus.batch_put([[1, 2, 3, 10, 11]]) + corpus.synchronize() + + ids, masks = _batch_get(corpus, [[1, 2, 3]]) + ids_list = ids.tolist() + leaf_paths = corpus.leaf_paths_from_mask(ids_list, masks.reshape(6, 6).tolist()) + self.assertIn([3, 10, 11], leaf_paths) + self.assertIn([3, 10, 99], leaf_paths) + self.assertEqual(ids_list.count(0), 2, ids_list) + + def test_external_sam_prob_prefers_frequent_continuation(self): + corpus = _make_corpus( + "PROB", + draft_token_num=2, + min_bfs_breadth=1, + max_bfs_breadth=1, + external_sam_budget=1, + external_corpus_documents=[ + [1, 2, 3, 10], + [1, 2, 3, 20], + [1, 2, 3, 20], + [1, 2, 3, 20], + ], + ) + + ids, _ = _batch_get(corpus, [[1, 2, 3]]) + self.assertEqual(ids.tolist(), [3, 20]) + + class TestNgramCorpusMatchBenchmark(CustomTestCase): """Benchmark incremental advance vs full rebuild in match()."""