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

[RadixTree] Reduce Stack Push/Pop Overhead for Leaf Nodes, Improve radix_tree Leaf Collection Performance (#12199)

Browse Source 
Signed-off-by: CLFutureX <chenyongqyl@163.com>
This commit is contained in:
PiteXChen
2025-11-11 23:35:39 +08:00
committed by GitHub
parent f09eee036d
commit 527b7d3f59
1 changed files with 29 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
python/sglang/srt/mem_cache/cpp_radix_tree/tree_v2_impl.h
View File

@@ -115,20 +115,24 @@ struct RadixTree::Impl {
std::vector<TreeNode*> collect_leaves() const {
std::vector<TreeNode*> leaves;
std::vector<TreeNode*> stack = {};
for (const auto& [_, child] : m_root) {
stack.push_back(child.get());
}
while (!stack.empty()) {
const auto node = stack.back();
stack.pop_back();
auto process_node = [&](TreeNode* node) {
if (node->is_leaf()) {
if (node->ref_count == 0) {
leaves.push_back(node);
}
} else {
for (const auto& [_, child] : *node) {
stack.push_back(child.get());
}
stack.push_back(node);
}
};
for (const auto& [_, child] : m_root) {
process_node(child.get());
}
while (!stack.empty()) {
const auto node = stack.back();
stack.pop_back();
for (const auto& [_, child] : *node) {
process_node(child.get());
}
}
return leaves;
@@ -139,23 +143,29 @@ struct RadixTree::Impl {
if (!use_hicache) return collect_leaves();
std::vector<TreeNode*> leaves;
std::vector<TreeNode*> stack = {};
for (const auto& [_, child] : m_root) {
stack.push_back(child.get());
}
while (!stack.empty()) {
const auto node = stack.back();
stack.pop_back();
if (!node->on_gpu()) continue; // skip nodes that are not on GPU
auto process_node = [&](TreeNode* node) {
if (!node->on_gpu()) return; // skip nodes that are not on GPU
if (node->is_leaf_device()) {
if (node->ref_count == 0) {
leaves.push_back(node);
}
} else {
for (const auto& [_, child] : *node) {
stack.push_back(child.get());
}
stack.push_back(node);
}
};
for (const auto& [_, child] : m_root) {
process_node(child.get())
}
while (!stack.empty()) {
const auto node = stack.back();
stack.pop_back();
for (const auto& [_, child] : *node) {
process_node(child.get());
}
}
return leaves;
}