chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Spec][Ngram] 1/N: Reference based Speculative Decoding refactor (#20393)

Browse Source
This commit is contained in:
kpham-sgl
2026-03-22 00:55:10 -07:00
committed by GitHub
parent d9f5c2179c
commit 6d160b42bb
13 changed files with 1039 additions and 388 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/advanced_features/server_arguments.md
View File

@@ -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 |

4
docs/advanced_features/speculative_decoding.md
View File

@@ -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 |

2
docs/platforms/ascend_npu_support_features.md
View File

@@ -228,7 +228,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
| `--speculative-ngram-`<br/>`max-match-window-size` | `12` | Type: int | Experimental |
| `--speculative-ngram-`<br/>`min-bfs-breadth` | `1` | Type: int | Experimental |
| `--speculative-ngram-`<br/>`max-bfs-breadth` | `10` | Type: int | Experimental |
| `--speculative-ngram-`<br/>`match-type` | `BFS` | `BFS`,<br/> `PROB` | Experimental |
| `--speculative-ngram-`<br/>`match-type` | `BFS` | `BFS`,<br/> `PROB` | Experimental. `BFS` uses recency-based expansion; `PROB` uses frequency-based expansion. |
| `--speculative-ngram-`<br/>`branch-length` | `18` | Type: int | Experimental |
| `--speculative-ngram-`<br/>`capacity` | `10000000` | Type: int | Experimental |

328
python/sglang/srt/speculative/cpp_ngram/ngram.cpp
View File

@@ -1,74 +1,16 @@
#include "ngram.h"
#include <algorithm>
#include <chrono>
#include <cstring>
#include <limits>
#include <list>
#include <mutex>
#include <queue>
#include <stdexcept>
#include <string>
#include <thread>
#include <tuple>
#include <unordered_map>
#include <vector>
#include "trie.h"
namespace ngram {
struct Node {
std::unordered_map<int32_t, int32_t> next;
};
Ngram::Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root) {
Ngram::Result info;
std::vector<int32_t> prevs;
info.token.reserve(draft_token_num);
prevs.reserve(draft_token_num);
std::queue<std::tuple<int32_t, int32_t, int32_t>> 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<Trie>(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<std::pair<TrieNode*, int32_t>> Ngram::match(const std::vector<int32_t>& 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<std::pair<TrieNode*, int32_t>> 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<int32_t> data;
if (!insert_queue_.dequeue(data)) {
continue;
}
const auto* token = data.data();
size_t size = data.size();
std::unique_lock<std::mutex> 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<std::vector<int32_t>>&& tokens) {
for (auto&& token : tokens) {
insert_queue_.enqueue(std::move(token));
}
}
Ngram::Result Ngram::matchBFS(const std::vector<int32_t>& tokens, size_t batch_size) const {
std::vector<std::pair<TrieNode*, int32_t>> 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<Node> tree(draft_token_num + 1);
int root = 0;
int cursor = 1;
for (auto [node, depth] : nodes) {
std::queue<std::tuple<int32_t, double, const TrieNode*>> 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<int32_t> data;
if (!insert_queue_.dequeue(data)) {
continue;
}
std::unique_lock<std::mutex> 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<int32_t>& tokens, size_t batch_size) const {
std::vector<std::pair<TrieNode*, int32_t>> nodes = match(tokens, batch_size);
auto draft_token_num = param_.get_draft_token_num(batch_size);
struct CompareByLastDouble {
bool operator()(
const std::tuple<double, const TrieNode*, double>& a, // parent_pos, node, final_prob
const std::tuple<double, const TrieNode*, double>& b) const {
return std::get<2>(a) < std::get<2>(b);
}
};
std::priority_queue<
std::tuple<double, const TrieNode*, double>,
std::vector<std::tuple<double, const TrieNode*, double>>,
CompareByLastDouble>
heap;
std::vector<Node> 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<std::pair<TrieNode*, int32_t>> topk_children;
for (auto* child : trie_node->sorted_children) {
sum_freq += static_cast<double>(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<double>(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<std::vector<int32_t>>& tokens) const {
Result Ngram::batchMatch(const std::vector<std::vector<int32_t>>& tokens) const {
std::unique_lock<std::mutex> 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

76
python/sglang/srt/speculative/cpp_ngram/ngram.h
View File

@@ -2,99 +2,42 @@
#include <cstddef>
#include <cstdint>
#include <functional>
#include <list>
#include <memory>
#include <mutex>
#include <new>
#include <set>
#include <sstream>
#include <thread>
#include <tuple>
#include <unordered_map>
#include <vector>
#include "param.h"
#include "queue.h"
#include "result.h"
#include "trie.h"
namespace ngram {
struct TrieNode {
std::unordered_map<int32_t, TrieNode*> child;
std::list<TrieNode*>::const_iterator global_lru_pos;
std::list<TrieNode*>::const_iterator parent_lru_pos;
int32_t token;
TrieNode* parent;
std::list<TrieNode*> 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<TrieNode*, CompareByFreq> sorted_children;
};
class Ngram {
std::vector<TrieNode> nodes_;
std::vector<TrieNode*> node_pool_;
size_t free_node_count_;
std::list<TrieNode*> global_lru_;
TrieNode* root_;
std::vector<TrieNode*> path_;
std::unique_ptr<Trie> trie_;
Param param_;
std::vector<std::pair<TrieNode*, int32_t>> match(const std::vector<int32_t>& 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<std::vector<int32_t>> insert_queue_;
std::thread insert_worker_;
std::vector<std::tuple<int32_t, int32_t, int32_t, int32_t>> 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<std::vector<int32_t>>&& tokens);
struct Result {
std::vector<int32_t> token;
std::vector<uint8_t> mask;
void truncate(size_t n);
};
Result batchMatch(const std::vector<std::vector<int32_t>>& tokens) const;
void reset() {
std::unique_lock<std::mutex> 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<int32_t>& tokens, size_t batch_size) const;
Result matchProb(const std::vector<int32_t>& tokens, size_t batch_size) const;
void insert();
void insertWorker();
};
} // namespace ngram

32
python/sglang/srt/speculative/cpp_ngram/ngram_cache.py → python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
View File

@@ -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)

10
python/sglang/srt/speculative/cpp_ngram/ngram_cache_binding.cpp → python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp
View File

@@ -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_<Ngram::Result>(m, "Result")
py::class_<Result>(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);
}

61
python/sglang/srt/speculative/cpp_ngram/result.cpp Normal file
View File

@@ -0,0 +1,61 @@
#include "result.h"
#include <cstring>
#include <queue>
#include <tuple>
namespace ngram {
Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root) {
Result info;
std::vector<int32_t> prevs;
info.token.reserve(draft_token_num);
prevs.reserve(draft_token_num);
std::queue<std::tuple<int32_t, int32_t, int32_t>> 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<size_t>(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

22
python/sglang/srt/speculative/cpp_ngram/result.h Normal file
View File

@@ -0,0 +1,22 @@
#pragma once
#include <cstdint>
#include <unordered_map>
#include <vector>
namespace ngram {
struct Result {
std::vector<int32_t> token;
std::vector<uint8_t> mask;
void truncate(size_t n);
};
struct Node {
std::unordered_map<int32_t, int32_t> next;
};
Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root);
} // namespace ngram

231
python/sglang/srt/speculative/cpp_ngram/trie.cpp Normal file
View File

@@ -0,0 +1,231 @@
#include "trie.h"
#include <algorithm>
#include <cstring>
#include <list>
#include <queue>
#include <tuple>
#include <vector>
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<size_t>(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<std::pair<TrieNode*, int32_t>>
Trie::match(const int32_t* context, size_t len, size_t min_window, size_t max_window) const {
std::vector<std::pair<TrieNode*, int32_t>> result;
result.reserve(max_window - min_window);
for (int32_t match_window_size = std::min(len, max_window); match_window_size >= static_cast<int32_t>(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<Node> tree(draft_token_num + 1);
int root = 0;
int cursor = 1;
for (auto [node, depth] : anchors) {
std::queue<std::tuple<int32_t, double, const TrieNode*>> queue;
queue.push({root, (param.max_match_window_size - depth) * bfs_breadth_scale + param.min_bfs_breadth, node});
while (queue.size() && cursor <= static_cast<int>(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<int>(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<double, const TrieNode*, double>& a,
const std::tuple<double, const TrieNode*, double>& b) const {
return std::get<2>(a) < std::get<2>(b);
}
};
std::priority_queue<
std::tuple<double, const TrieNode*, double>,
std::vector<std::tuple<double, const TrieNode*, double>>,
CompareByLastDouble>
heap;
std::vector<Node> 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<std::pair<TrieNode*, int32_t>> topk_children;
for (auto* child : trie_node->sorted_children) {
sum_freq += static_cast<double>(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<double>(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<int>(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

71
python/sglang/srt/speculative/cpp_ngram/trie.h Normal file
View File

@@ -0,0 +1,71 @@
#pragma once
#include <cstddef>
#include <cstdint>
#include <functional>
#include <list>
#include <new>
#include <set>
#include <tuple>
#include <unordered_map>
#include <vector>
#include "param.h"
#include "result.h"
namespace ngram {
struct TrieNode {
std::unordered_map<int32_t, TrieNode*> child;
std::list<TrieNode*>::const_iterator global_lru_pos;
std::list<TrieNode*>::const_iterator parent_lru_pos;
int32_t token;
TrieNode* parent;
std::list<TrieNode*> 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<TrieNode*, CompareByFreq> 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<std::pair<TrieNode*, int32_t>>
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<TrieNode> nodes_;
std::vector<TrieNode*> node_pool_;
size_t free_node_count_;
std::list<TrieNode*> global_lru_;
TrieNode* root_;
std::vector<TrieNode*> path_;
Param param_;
};
} // namespace ngram

17
python/sglang/srt/speculative/ngram_worker.py
View File

@@ -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:

571
test/registered/spec/utils/test_ngram_corpus.py Normal file
View File

@@ -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)