From 527b7d3f596cb102fc0a90fe38372649bd70fe42 Mon Sep 17 00:00:00 2001
From: PiteXChen <44110731+CLFutureX@users.noreply.github.com>
Date: Tue, 11 Nov 2025 23:35:39 +0800
Subject: [PATCH] [RadixTree] Reduce Stack Push/Pop Overhead for Leaf Nodes,
 Improve radix_tree Leaf Collection Performance  (#12199)

Signed-off-by: CLFutureX <chenyongqyl@163.com>
---
 .../mem_cache/cpp_radix_tree/tree_v2_impl.h   | 48 +++++++++++--------
 1 file changed, 29 insertions(+), 19 deletions(-)

diff --git a/python/sglang/srt/mem_cache/cpp_radix_tree/tree_v2_impl.h b/python/sglang/srt/mem_cache/cpp_radix_tree/tree_v2_impl.h
index cb9f9dde5..c674d2722 100644
--- a/python/sglang/srt/mem_cache/cpp_radix_tree/tree_v2_impl.h
+++ b/python/sglang/srt/mem_cache/cpp_radix_tree/tree_v2_impl.h
@@ -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;
   }