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

[Spec][Ngram] 2/N: Rename branch length to max trie depth (#21181)

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
kpham-sgl
2026-03-22 23:35:25 -07:00
committed by GitHub
parent d6b12c401c
commit bc4aaab6a1
11 changed files with 31 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/server_args.py
View File

@@ -506,7 +506,7 @@ class ServerArgs:
speculative_ngram_min_bfs_breadth: int = 1
speculative_ngram_max_bfs_breadth: int = 10
speculative_ngram_match_type: Literal["BFS", "PROB"] = "BFS"
speculative_ngram_branch_length: int = 18
speculative_ngram_max_trie_depth: int = 18
speculative_ngram_capacity: int = 10 * 1000 * 1000
enable_multi_layer_eagle: bool = False
@@ -4766,10 +4766,10 @@ class ServerArgs:
help="The match type for cache tree.",
)
parser.add_argument(
"--speculative-ngram-branch-length",
"--speculative-ngram-max-trie-depth",
type=int,
default=ServerArgs.speculative_ngram_branch_length,
help="The branch length for ngram speculative decoding.",
default=ServerArgs.speculative_ngram_max_trie_depth,
help="The max trie depth for ngram speculative decoding.",
)
parser.add_argument(
"--speculative-ngram-capacity",

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

@@ -11,9 +11,9 @@
namespace ngram {
Ngram::Ngram(size_t capacity, const Param& param) : param_(param) {
if (!(param_.branch_length > 1)) {
if (!(param_.max_trie_depth > 1)) {
throw std::runtime_error(
"param_.branch_length must be greater than 1, current value: " + std::to_string(param_.branch_length));
"param_.max_trie_depth must be greater than 1, current value: " + std::to_string(param_.max_trie_depth));
}
if (!(param_.min_match_window_size > 0)) {
throw std::runtime_error(
@@ -26,11 +26,11 @@ Ngram::Ngram(size_t capacity, const Param& param) : param_(param) {
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)) {
if (!(param_.max_match_window_size < param_.max_trie_depth)) {
throw std::runtime_error(
"max_match_window_size must be less than branch_length, current "
"max_match_window_size must be less than max_trie_depth, current "
"max_match_window_size: " +
std::to_string(param_.max_match_window_size) + ", branch_length: " + std::to_string(param_.branch_length));
std::to_string(param_.max_match_window_size) + ", max_trie_depth: " + std::to_string(param_.max_trie_depth));
}
if (!(param_.min_bfs_breadth > 0)) {
throw std::runtime_error(

6
python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
View File

@@ -25,7 +25,7 @@ ngram_corpus_cpp = load(
class NgramCorpus:
def __init__(
self,
branch_length=18,
max_trie_depth=18,
min_match_window_size=1,
max_match_window_size=10,
min_bfs_breadth=1,
@@ -35,7 +35,7 @@ class NgramCorpus:
capacity=1000000,
):
param = ngram_corpus_cpp.Param()
param.branch_length = branch_length
param.max_trie_depth = max_trie_depth
param.min_match_window_size = min_match_window_size
param.max_match_window_size = max_match_window_size
param.min_bfs_breadth = min_bfs_breadth
@@ -131,7 +131,7 @@ if __name__ == "__main__":
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
[1, 2, 3, 44, 55, 66, 77, 88, 99, 100],
]
corpus = NgramCorpus(branch_length=12, draft_token_num=8)
corpus = NgramCorpus(max_trie_depth=12, draft_token_num=8)
corpus.batch_put(token_ids)
corpus.synchronize()

2
python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp
View File

@@ -23,7 +23,7 @@ PYBIND11_MODULE(ngram_corpus_cpp, m) {
.def_readwrite("max_bfs_breadth", &Param::max_bfs_breadth)
.def_readwrite("min_match_window_size", &Param::min_match_window_size)
.def_readwrite("max_match_window_size", &Param::max_match_window_size)
.def_readwrite("branch_length", &Param::branch_length)
.def_readwrite("max_trie_depth", &Param::max_trie_depth)
.def_readwrite("draft_token_num", &Param::draft_token_num)
.def_readwrite("match_type", &Param::match_type)
.def_readwrite("batch_min_match_window_size", &Param::batch_min_match_window_size)

4
python/sglang/srt/speculative/cpp_ngram/param.h
View File

@@ -19,7 +19,7 @@ struct Param {
size_t max_bfs_breadth;
size_t min_match_window_size;
size_t max_match_window_size;
size_t branch_length;
size_t max_trie_depth;
size_t draft_token_num;
std::string match_type;
@@ -109,7 +109,7 @@ struct Param {
ss << "enable = " << enable << ", enable_router_mode = " << enable_router_mode
<< ", min_bfs_breadth = " << min_bfs_breadth << ", max_bfs_breadth = " << max_bfs_breadth
<< ", min_match_window_size = " << min_match_window_size << ", max_match_window_size = " << max_match_window_size
<< ", branch_length = " << branch_length << ", draft_token_num = " << draft_token_num
<< ", max_trie_depth = " << max_trie_depth << ", draft_token_num = " << draft_token_num
<< ", match_type = " << match_type;
ss << ", batch_min_match_window_size(" << batch_min_match_window_size.size() << ") = ";
for (int i = 0; i < batch_min_match_window_size.size(); ++i) {

2
python/sglang/srt/speculative/cpp_ngram/trie.cpp
View File

@@ -21,7 +21,7 @@ Trie::Trie(size_t capacity, const Param& param) : param_(param) {
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);
auto end = start + std::min(len - i, param_.max_trie_depth);
if (static_cast<size_t>(end - start) > free_node_count_) {
squeeze(end - start - free_node_count_);

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

@@ -40,7 +40,7 @@ class NGRAMWorker:
self.tp_rank = tp_rank
self.page_size = server_args.page_size
self.draft_token_num: int = server_args.speculative_num_draft_tokens
self.branch_length: int = server_args.speculative_ngram_branch_length
self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth
self.max_match_window_size: int = (
server_args.speculative_ngram_max_match_window_size
)
@@ -57,7 +57,7 @@ class NGRAMWorker:
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,
max_trie_depth=server_args.speculative_ngram_max_trie_depth,
draft_token_num=server_args.speculative_num_draft_tokens,
)
@@ -209,7 +209,7 @@ class NGRAMWorker:
# put_ids = req.origin_input_ids + req.output_ids
# else:
put_ids = self._efficient_concat_last_n(
req.origin_input_ids, req.output_ids, self.branch_length
req.origin_input_ids, req.output_ids, self.max_trie_depth
)
batch_tokens.append(put_ids)
self.ngram_corpus.batch_put(batch_tokens)