minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Radix Cache Split: Spin off TreeCore (#29901)

Browse Source
This commit is contained in:
Jialin Ouyang
2026-07-25 14:31:59 -07:00
committed by GitHub
parent a23f6ea090
commit cd145f840f
25 changed files with 6697 additions and 2703 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+3 -1
View File
@@ -1249,7 +1249,9 @@ class Req(ReqDllmMixin):
)
)
if envs.SGLANG_RADIX_FORCE_MISS.get():
match_result = zero_match_result(tree_cache, match_result)
match_result = zero_match_result(
tree_cache, match_result, extra_key=self.extra_key
)
(
self.prefix_indices,
self.last_node,
python/sglang/srt/managers/schedule_policy.py
+6 -3
View File
@@ -115,7 +115,9 @@ def match_prefix_for_req(
)
)
if envs.SGLANG_RADIX_FORCE_MISS.get():
match_result = zero_match_result(tree_cache, match_result)
match_result = zero_match_result(
tree_cache, match_result, extra_key=req.extra_key
)
(
req.prefix_indices,
req.last_node,
@@ -290,7 +292,7 @@ class SchedulePolicy:
)
if envs.SGLANG_RADIX_FORCE_MISS.get():
match_result = zero_match_result(
self.waiting_queue_radix_tree, match_result
self.waiting_queue_radix_tree, match_result, extra_key=extra_key
)
in_batch_matching_prefixes = match_result.device_indices
if (
@@ -328,7 +330,8 @@ class SchedulePolicy:
"""Sorts the waiting queue based on a depth-first search weighting."""
last_node_to_reqs = defaultdict(list)
for req in waiting_queue:
last_node_to_reqs[req.last_node].append(req)
last_node = tree_cache.resolve_node_handle(req.last_node)
last_node_to_reqs[last_node].append(req)
node_to_weight = defaultdict(int)
for node in last_node_to_reqs:
python/sglang/srt/managers/scheduler.py
+9 -9
View File
@@ -2418,25 +2418,25 @@ class Scheduler(
def _prefetch_kvcache(self, req: Req):
if self.enable_hicache_storage:
req.init_next_round_input(self.tree_cache, cow_mamba=False)
last_host_node = req.last_host_node
if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
last_hash = last_host_node.get_last_hash_value()
tree_cache = self.tree_cache
if tree_cache.is_backuped(req.last_host_node) or tree_cache.is_root(
req.last_host_node
):
matched_len = len(req.prefix_indices) + req.host_hit_length
match_end = req._compute_max_prefix_len(
len(req.full_untruncated_fill_ids)
)
new_input_tokens = req.full_untruncated_fill_ids[matched_len:match_end]
prefix_keys = (
last_host_node.get_prefix_hash_values(last_host_node.parent)
if self.tree_cache.hicache_storage_pass_prefix_keys
tree_cache.get_prefix_hash_values(req.last_host_node)
if tree_cache.hicache_storage_pass_prefix_keys
else None
)
self.tree_cache.prefetch_from_storage(
tree_cache.prefetch_from_storage(
req.rid,
last_host_node,
req.last_host_node,
new_input_tokens,
last_hash,
tree_cache.get_last_hash_value(req.last_host_node),
prefix_keys,
)