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

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: