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

[RadixTree][8/N Refactor]: unify lock interface (#20330)

Browse Source
This commit is contained in:
pansicheng
2026-03-16 11:49:51 +08:00
committed by GitHub
parent f0458e0b49
commit 97b2a89334
11 changed files with 144 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/managers/schedule_policy.py
View File

@@ -3,6 +3,7 @@ from __future__ import annotations
import logging
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
from sglang.srt.utils import get_bool_env_var
_ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG")
@@ -560,11 +561,9 @@ class PrefillAdder:
self._update_prefill_budget(prefix_len, trunc_len, 0)
def _req_inc_lock_ref(self, req: Req):
result = self.tree_cache.inc_lock_ref(req.last_node)
if self.is_hybrid_swa:
swa_uuid_for_lock = self.tree_cache.inc_lock_ref(req.last_node)
req.swa_uuid_for_lock = swa_uuid_for_lock
else:
self.tree_cache.inc_lock_ref(req.last_node)
req.swa_uuid_for_lock = result.swa_uuid_for_lock
def add_dllm_staging_req(self, req: Req):
assert self.dllm_config is not None
@@ -624,14 +623,15 @@ class PrefillAdder:
@contextmanager
def _lock_node(self, last_node: TreeNode):
try:
result = self.tree_cache.inc_lock_ref(last_node)
if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
swa_uuid_for_lock = self.tree_cache.inc_lock_ref(last_node)
else:
self.tree_cache.inc_lock_ref(last_node)
swa_uuid_for_lock = result.swa_uuid_for_lock
yield None
finally:
if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock)
self.tree_cache.dec_lock_ref(
last_node, DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock)
)
else:
self.tree_cache.dec_lock_ref(last_node)

28
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -88,6 +88,28 @@ class EvictResult:
mamba_num_evicted: int = 0
@dataclasses.dataclass
class IncLockRefResult:
"""Result of an inc_lock_ref operation."""
delta: Optional[int] = None
swa_uuid_for_lock: Optional[int] = None
@dataclasses.dataclass
class DecLockRefParams:
"""Parameters for dec_lock_ref operation."""
swa_uuid_for_lock: Optional[int] = None
@dataclasses.dataclass
class DecLockRefResult:
"""Result of an dec_lock_ref operation."""
delta: Optional[int] = None
class MatchResult(NamedTuple):
"""Result of a prefix match operation.
@@ -158,11 +180,13 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
pass
@abstractmethod
def inc_lock_ref(self, node: Any):
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
pass
@abstractmethod
def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
def dec_lock_ref(
self, node: Any, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
pass
def evictable_size(self):

13
python/sglang/srt/mem_cache/chunk_cache.py
View File

@@ -9,8 +9,11 @@ import torch
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InsertParams,
InsertResult,
MatchPrefixParams,
@@ -80,11 +83,13 @@ class ChunkCache(BasePrefixCache):
def evict(self, params: EvictParams) -> EvictResult:
return EvictResult()
def inc_lock_ref(self, node: Any):
return 0
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
return IncLockRefResult(delta=0)
def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
return 0
def dec_lock_ref(
self, node: Any, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
return DecLockRefResult(delta=0)
def protected_size(self):
# NOTE: no protected size in chunk cache. Chunk cache's eviction is the same with request's lifecycle.

20
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
View File

@@ -14,8 +14,11 @@ import torch
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
MatchPrefixParams,
MatchResult,
)
@@ -198,7 +201,8 @@ class HiMambaRadixCache(MambaRadixCache):
else:
ancestor_node = node
delta = self.inc_lock_ref(ancestor_node)
result = self.inc_lock_ref(ancestor_node)
delta = result.delta
full_host_indices = torch.cat([n.host_value for n in nodes_to_load])
if (len(full_host_indices) < self.load_back_threshold) or (
@@ -988,9 +992,9 @@ class HiMambaRadixCache(MambaRadixCache):
return
super().sanity_check()
def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
if self.disable:
return 0
return IncLockRefResult(delta=0)
delta = 0
if node.mamba_value is not None:
@@ -1014,11 +1018,13 @@ class HiMambaRadixCache(MambaRadixCache):
self.evictable_full_device_leaves.discard(node)
node.full_lock_ref += 1
node = node.parent
return delta
return IncLockRefResult(delta=delta)
def dec_lock_ref(self, node: TreeNode):
def dec_lock_ref(
self, node: TreeNode, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
if self.disable:
return 0
return DecLockRefResult(delta=0)
delta = 0
@@ -1044,7 +1050,7 @@ class HiMambaRadixCache(MambaRadixCache):
if node.full_lock_ref == 0:
self._update_full_device_leaf_status(node)
node = node.parent
return delta
return DecLockRefResult(delta=delta)
# ---- L3 Support ----

20
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -15,8 +15,11 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InsertParams,
InsertResult,
MatchPrefixParams,
@@ -817,9 +820,9 @@ class HiRadixCache(RadixCache):
"""
return RadixKey(token_ids=list(token_ids))
def inc_lock_ref(self, node: TreeNode):
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
if self.disable:
return 0
return IncLockRefResult(delta=0)
delta = 0
while node != self.root_node:
@@ -831,11 +834,13 @@ class HiRadixCache(RadixCache):
self._update_leaf_status(node)
self._update_host_leaf_status(node)
node = node.parent
return delta
return IncLockRefResult(delta=delta)
def dec_lock_ref(self, node: TreeNode):
def dec_lock_ref(
self, node: TreeNode, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
if self.disable:
return 0
return DecLockRefResult(delta=0)
delta = 0
while node != self.root_node:
@@ -851,7 +856,7 @@ class HiRadixCache(RadixCache):
node is self.root_node
), f"This request holds the node from another tree"
node = node.parent
return delta
return DecLockRefResult(delta=delta)
def _update_host_leaf_status(self, node: TreeNode):
if not node.evicted or node.lock_ref > 0:
@@ -1011,7 +1016,8 @@ class HiRadixCache(RadixCache):
ancester_node = node
# protect the ancestor nodes from eviction
delta = self.inc_lock_ref(ancester_node)
result = self.inc_lock_ref(ancester_node)
delta = result.delta
# load it all or not at all
host_indices = torch.cat([n.host_value for n in nodes_to_load])

17
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -35,8 +35,11 @@ from sglang.srt.mem_cache.allocator import (
)
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InsertParams,
InsertResult,
MatchPrefixParams,
@@ -806,14 +809,14 @@ class MambaRadixCache(BasePrefixCache):
return full_num_evicted
def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
"""
Increment the lock reference count for the node.
It locks the full_lock_ref for nodes between the [last node, root), exclusive.
It locks the mamba_lock_ref for current node if its mamba_value exists.
"""
if self.disable:
return None
return IncLockRefResult()
# protect mamba value in current node if it exists
if node.mamba_value is not None:
@@ -832,16 +835,18 @@ class MambaRadixCache(BasePrefixCache):
self.full_protected_size_ += len(node.value)
node.full_lock_ref += 1
node = node.parent
return None
return IncLockRefResult()
def dec_lock_ref(self, node: TreeNode):
def dec_lock_ref(
self, node: TreeNode, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
"""
Decrement the lock reference count for the node.
It unlocks the full_lock_ref for nodes between the [last node, root), exclusive.
It unlocks the mamba_lock_ref for current node if its mamba_value exists.
"""
if self.disable:
return None
return DecLockRefResult()
if node.mamba_value is not None:
assert (
@@ -862,6 +867,8 @@ class MambaRadixCache(BasePrefixCache):
node.full_lock_ref -= 1
node = node.parent
return DecLockRefResult()
def sanity_check(self):
if self.disable:
return

17
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -42,8 +42,11 @@ from sglang.srt.disaggregation.kv_events import (
)
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InsertParams,
InsertResult,
MatchPrefixParams,
@@ -609,9 +612,9 @@ class RadixCache(BasePrefixCache):
self.update_eviction_metrics(num_evicted, start_time)
return EvictResult(num_tokens_evicted=num_evicted)
def inc_lock_ref(self, node: TreeNode):
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
if self.disable:
return 0
return IncLockRefResult(delta=0)
delta = 0
while node != self.root_node:
@@ -622,11 +625,13 @@ class RadixCache(BasePrefixCache):
node.lock_ref += 1
self._update_leaf_status(node)
node = node.parent
return delta
return IncLockRefResult(delta=delta)
def dec_lock_ref(self, node: TreeNode):
def dec_lock_ref(
self, node: TreeNode, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
if self.disable:
return 0
return DecLockRefResult(delta=0)
delta = 0
while node != self.root_node:
@@ -641,7 +646,7 @@ class RadixCache(BasePrefixCache):
node is self.root_node
), f"This request holds the node from another tree"
node = node.parent
return delta
return DecLockRefResult(delta=delta)
def evictable_size(self):
return self.evictable_size_

13
python/sglang/srt/mem_cache/radix_cache_cpp.py
View File

@@ -2,14 +2,17 @@ from __future__ import annotations
import logging
import time
from typing import TYPE_CHECKING, List, Set
from typing import TYPE_CHECKING, List, Optional, Set
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
MatchPrefixParams,
MatchResult,
)
@@ -123,21 +126,25 @@ class RadixCacheCpp(BasePrefixCache):
raise NotImplementedError("Host cache is not supported yet")
def dec_lock_ref(self, node: TreeNodeCpp):
def dec_lock_ref(
self, node: TreeNodeCpp, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
"""
Decrement the reference count of a node to root of the radix tree.
Args:
node (TreeNodeCpp): The handle of the node to decrement the reference count for.
"""
self.tree.lock_ref(node, False) # do not increment
return DecLockRefResult()
def inc_lock_ref(self, node: TreeNodeCpp):
def inc_lock_ref(self, node: TreeNodeCpp) -> IncLockRefResult:
"""
Increment the reference count of from a node to root of the radix tree.
Args:
node (TreeNodeCpp): The handle of the node to increment the reference count for.
"""
self.tree.lock_ref(node, True)
return IncLockRefResult()
def evict(self, params: EvictParams) -> EvictResult:
start_time = time.perf_counter()

22
python/sglang/srt/mem_cache/session_aware_cache.py
View File

@@ -7,8 +7,11 @@ import torch
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
MatchPrefixParams,
MatchResult,
)
@@ -223,17 +226,17 @@ class SessionAwareCache(BasePrefixCache):
def evict(self, params: EvictParams) -> EvictResult:
return self.inner.evict(params)
def inc_lock_ref(self, node: Any):
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
if isinstance(node, _VirtualNode):
return None
return IncLockRefResult()
return self.inner.inc_lock_ref(node)
def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
def dec_lock_ref(
self, node: Any, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
if isinstance(node, _VirtualNode):
return
if swa_uuid_for_lock is not None:
return self.inner.dec_lock_ref(node, swa_uuid_for_lock)
return self.inner.dec_lock_ref(node)
return DecLockRefResult()
return self.inner.dec_lock_ref(node, params)
# -- Session lifecycle --
@@ -245,7 +248,10 @@ class SessionAwareCache(BasePrefixCache):
if slot.last_node is not None:
if slot.swa_uuid_for_lock is not None:
self.inner.dec_lock_ref(slot.last_node, slot.swa_uuid_for_lock)
self.inner.dec_lock_ref(
slot.last_node,
DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock),
)
else:
self.inner.dec_lock_ref(slot.last_node)

30
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -30,8 +30,11 @@ from numpy import float64
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
DecLockRefResult,
EvictParams,
EvictResult,
IncLockRefResult,
InsertParams,
InsertResult,
MatchPrefixParams,
@@ -485,7 +488,9 @@ class SWARadixCache(BasePrefixCache):
self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
# Remove req slot release the cache lock
self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock)
self.dec_lock_ref(
req.last_node, DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock)
)
def cache_unfinished_req(self, req: Req, chunked=False) -> None:
"""Cache request when it is unfinished."""
@@ -536,8 +541,11 @@ class SWARadixCache(BasePrefixCache):
req.cache_protected_len = len(new_indices)
self.dec_lock_ref(req.last_node, req.swa_uuid_for_lock)
swa_uuid_for_lock = self.inc_lock_ref(new_last_node)
self.dec_lock_ref(
req.last_node, DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock)
)
result = self.inc_lock_ref(new_last_node)
swa_uuid_for_lock = result.swa_uuid_for_lock
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
if len(new_indices) < len(kv_indices):
@@ -647,7 +655,7 @@ class SWARadixCache(BasePrefixCache):
num_tokens_evicted=full_num_evicted, swa_num_tokens_evicted=swa_num_evicted
)
def inc_lock_ref(self, node: TreeNode) -> Optional[int]:
def inc_lock_ref(self, node: TreeNode) -> IncLockRefResult:
"""
Increment the lock reference count for the node. Returns the swa_uuid_for_lock, which needs
to be passed to dec_lock_ref.
@@ -655,7 +663,7 @@ class SWARadixCache(BasePrefixCache):
It locks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive.
"""
if self.disable:
return None
return IncLockRefResult()
swa_lock_size = 0
swa_uuid_for_lock = None
@@ -686,17 +694,21 @@ class SWARadixCache(BasePrefixCache):
node.swa_uuid = gen_swa_uuid()
swa_uuid_for_lock = node.swa_uuid
node = node.parent
return swa_uuid_for_lock
return IncLockRefResult(swa_uuid_for_lock=swa_uuid_for_lock)
def dec_lock_ref(self, node: TreeNode, swa_uuid_for_lock: Optional[int] = None):
def dec_lock_ref(
self, node: TreeNode, params: Optional[DecLockRefParams] = None
) -> DecLockRefResult:
"""
Decrement the lock reference count for the node.
It unlocks the full_lock_ref for nodes between the [last node, root), exclusive.
It unlocks the swa_lock_ref for nodes between the [last node, swa_uuid_for_lock], inclusive.
If swa_uuid_for_lock is None, it unlocks to the root, exclusive.
"""
swa_uuid_for_lock = params.swa_uuid_for_lock if params is not None else None
if self.disable:
return
return DecLockRefResult()
dec_lock_swa = True
while node != self.root_node:
@@ -725,6 +737,8 @@ class SWARadixCache(BasePrefixCache):
node = node.parent
return DecLockRefResult()
def sanity_check(self):
self.full_lru_list.sanity_check(self)
self.swa_lru_list.sanity_check(self)