diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md
index 961d45505..61cfe91e0 100644
--- a/docs/advanced_features/server_arguments.md
+++ b/docs/advanced_features/server_arguments.md
@@ -299,7 +299,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--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` | 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-max-trie-depth` | The max trie depth for ngram speculative decoding. | `18` | Type: int |
| `--speculative-ngram-capacity` | The cache capacity for ngram speculative decoding. | `10000000` | Type: int |
## Multi-layer Eagle speculative decoding
diff --git a/docs/advanced_features/speculative_decoding.md b/docs/advanced_features/speculative_decoding.md
index 0ea688d00..c573af072 100644
--- a/docs/advanced_features/speculative_decoding.md
+++ b/docs/advanced_features/speculative_decoding.md
@@ -393,7 +393,7 @@ Enable it with:
| `--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"` |
-| `--speculative-ngram-branch-length` | How many recent tokens to insert into the cache. | `18` |
+| `--speculative-ngram-max-trie-depth` | The max trie depth for ngram speculative decoding. | `18` |
| `--speculative-ngram-capacity` | Cache capacity (number of entries). | `10,000,000` |
Notes:
@@ -469,7 +469,7 @@ Below is a comprehensive list of all speculative decoding parameters available i
| `--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"` | 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-max-trie-depth` | `int` | `18` | Max trie depth for ngram speculative decoding |
| `--speculative-ngram-capacity` | `int` | `10,000,000` | Cache capacity |
### Environment variables
diff --git a/docs/platforms/ascend_npu_support_features.md b/docs/platforms/ascend_npu_support_features.md
index 4a6e9eeb5..ee243b798 100644
--- a/docs/platforms/ascend_npu_support_features.md
+++ b/docs/platforms/ascend_npu_support_features.md
@@ -229,7 +229,7 @@ click [Server Arguments](https://docs.sglang.io/advanced_features/server_argumen
| `--speculative-ngram-`
`min-bfs-breadth` | `1` | Type: int | Experimental |
| `--speculative-ngram-`
`max-bfs-breadth` | `10` | Type: int | Experimental |
| `--speculative-ngram-`
`match-type` | `BFS` | `BFS`,
`PROB` | Experimental. `BFS` uses recency-based expansion; `PROB` uses frequency-based expansion. |
-| `--speculative-ngram-`
`branch-length` | `18` | Type: int | Experimental |
+| `--speculative-ngram-`
`max-trie-depth` | `18` | Type: int | Experimental |
| `--speculative-ngram-`
`capacity` | `10000000` | Type: int | Experimental |
## Expert parallelism
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 3a8d7fbf8..f94da569a 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -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",
diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
index a5ca8c840..582c87659 100644
--- a/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
+++ b/python/sglang/srt/speculative/cpp_ngram/ngram.cpp
@@ -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(
diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
index 56d417bd4..e44a3da6b 100644
--- a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
+++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus.py
@@ -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()
diff --git a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp
index edc0e2c5c..8da395440 100644
--- a/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp
+++ b/python/sglang/srt/speculative/cpp_ngram/ngram_corpus_binding.cpp
@@ -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)
diff --git a/python/sglang/srt/speculative/cpp_ngram/param.h b/python/sglang/srt/speculative/cpp_ngram/param.h
index 08b975bb1..d31af64ba 100644
--- a/python/sglang/srt/speculative/cpp_ngram/param.h
+++ b/python/sglang/srt/speculative/cpp_ngram/param.h
@@ -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) {
diff --git a/python/sglang/srt/speculative/cpp_ngram/trie.cpp b/python/sglang/srt/speculative/cpp_ngram/trie.cpp
index 7c7340d57..8d9eec82b 100644
--- a/python/sglang/srt/speculative/cpp_ngram/trie.cpp
+++ b/python/sglang/srt/speculative/cpp_ngram/trie.cpp
@@ -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(end - start) > free_node_count_) {
squeeze(end - start - free_node_count_);
diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py
index 784a4130e..04a38cefb 100644
--- a/python/sglang/srt/speculative/ngram_worker.py
+++ b/python/sglang/srt/speculative/ngram_worker.py
@@ -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)
diff --git a/test/registered/spec/utils/test_ngram_corpus.py b/test/registered/spec/utils/test_ngram_corpus.py
index de1b20168..b34895885 100644
--- a/test/registered/spec/utils/test_ngram_corpus.py
+++ b/test/registered/spec/utils/test_ngram_corpus.py
@@ -11,7 +11,7 @@ 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,
+ max_trie_depth=12,
min_match_window_size=1,
max_match_window_size=10,
min_bfs_breadth=1,
@@ -240,7 +240,7 @@ class TestNgramCorpusSqueeze(CustomTestCase):
def test_eviction_preserves_recent(self):
corpus = _make_corpus(
- "BFS", capacity=500, branch_length=6, max_match_window_size=5
+ "BFS", capacity=500, max_trie_depth=6, max_match_window_size=5
)
old_seq = list(range(1000, 1050))
@@ -358,7 +358,7 @@ class TestFrequencyBoosting(CustomTestCase):
max_bfs_breadth=1,
min_bfs_breadth=1,
max_match_window_size=3,
- branch_length=5,
+ max_trie_depth=5,
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
@@ -387,7 +387,7 @@ class TestRecencyOrdering(CustomTestCase):
max_bfs_breadth=1,
min_bfs_breadth=1,
max_match_window_size=3,
- branch_length=5,
+ max_trie_depth=5,
)
corpus.batch_put([[1, 2, 3, 10, 11]])
corpus.synchronize()
@@ -433,10 +433,10 @@ class TestSingleTokenContext(CustomTestCase):
class TestLongContext(CustomTestCase):
- """Verify behavior when query context exceeds branch_length."""
+ """Verify behavior when query context exceeds max_trie_depth."""
- def test_context_longer_than_branch_length(self):
- corpus = _make_corpus("BFS", branch_length=6, max_match_window_size=5)
+ def test_context_longer_than_max_trie_depth(self):
+ corpus = _make_corpus("BFS", max_trie_depth=6, max_match_window_size=5)
seq = list(range(1, 20))
corpus.batch_put([seq])
corpus.synchronize()
@@ -539,7 +539,7 @@ class TestSqueezeEvictsOld(CustomTestCase):
def test_old_data_evicted(self):
corpus = _make_corpus(
- "BFS", capacity=150, branch_length=6, max_match_window_size=5
+ "BFS", capacity=150, max_trie_depth=6, max_match_window_size=5
)
old_seq = list(range(5000, 5030))