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

[UnifiedTree]: Fix UnifiedRadixCache device match semantics with HiCache (#25277)

Browse Source
This commit is contained in:
Zhangheng
2026-05-16 00:40:55 +08:00
committed by GitHub
parent 3f7e538b2f
commit 21b3ac52b4
7 changed files with 398 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache_components/full_component.py
+9 -2
View File
@@ -41,8 +41,15 @@ class FullComponent(TreeComponent):
# HiCache state: set to host KV pool when HiCache enabled
self._full_kv_pool_host = None
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
# HiCache: evicted + backuped nodes are valid match boundaries
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
if match_device_only:
return (
lambda node: node.component_data[self.component_type].value is not None
)
# HiCache: evicted + backuped nodes are valid match boundaries.
return lambda node: (
node.component_data[self.component_type].value is not None or node.backuped
)
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+10 -2
View File
@@ -50,8 +50,13 @@ class MambaComponent(TreeComponent):
# HiCache state
self._mamba_pool_host = None # set to host mamba pool when HiCache enabled
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
ct = self.component_type
if match_device_only:
return lambda node: node.component_data[ct].value is not None
# HiCache: evicted + backuped (host_value present) is also a valid match
return lambda node: (
node.component_data[ct].value is not None
@@ -69,7 +74,10 @@ class MambaComponent(TreeComponent):
req = params.req
last_node = result.best_match_node
if len(value_chunks) > best_value_len:
# HiCache can still use prefix matches and load back host-backed Mamba
# states. We temporarily skip branching-state fill in that mode and can
# add a HiCache-aware branching policy later.
if self.cache.cache_controller is None and len(value_chunks) > best_value_len:
chunk_size = get_global_server_args().mamba_cache_chunk_size
aligned_seqlen = (
sum(len(v) for v in value_chunks) // chunk_size
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+4 -2
View File
@@ -69,7 +69,9 @@ class SWAComponent(TreeComponent):
self.cache.lru_lists[ct].insert_mru(node)
self.cache.component_evictable_size_[ct] += len(value)
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
sliding_window_size = self.sliding_window_size
ct = self.component_type
state = {"len": float("inf")}
@@ -78,7 +80,7 @@ class SWAComponent(TreeComponent):
cd = node.component_data[ct]
# HiCache: a host-only tombstone is a valid match boundary too
# — load_back will restore SWA from host before use.
if cd.value is None and cd.host_value is None:
if cd.value is None and (match_device_only or cd.host_value is None):
state["len"] = 0
return False
state["len"] += len(node.key)
python/sglang/srt/mem_cache/unified_cache_components/tree_component.py
+6 -2
View File
@@ -116,11 +116,15 @@ class TreeComponent(ABC):
return len(value) if value is not None else 0
@abstractmethod
def create_match_validator(self) -> Callable[[UnifiedTreeNode], bool]:
def create_match_validator(
self, match_device_only: bool = False
) -> Callable[[UnifiedTreeNode], bool]:
"""Return a per-match stateful predicate that decides whether a node
is a valid match boundary for this component.
Called once per match_prefix; the returned closure may carry state.
- Full: always True (every node is valid).
When match_device_only is true, host-backed nodes must not be accepted
as valid match boundaries.
- Full: returns True if the node has full component data.
- SWA: tracks accumulated length since last gap; returns True only
when the contiguous window reaches swa_sliding_window_size.
- Mamba: returns True iff the node has mamba component data."""
python/sglang/srt/mem_cache/unified_radix_cache.py
+100 -84
View File
@@ -351,9 +351,18 @@ class UnifiedRadixCache(BasePrefixCache):
if len(key) == 0:
return self._empty_match_result
value, best_match_node, best_value_len = self._match_prefix_helper(key)
(
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
) = self._match_prefix_helper(key)
return self._match_post_processor(
params, value, best_match_node, best_value_len
params,
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
)
def insert(self, params: InsertParams) -> InsertResult:
@@ -585,67 +594,53 @@ class UnifiedRadixCache(BasePrefixCache):
# ---- Internal Helpers ----
def _match_prefix_helper_readonly(
self, key: RadixKey
) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
"""Read-only version of _match_prefix_helper that does not split nodes.
Only considers fully matched nodes, ignores partial matches.
Not used yet; reserved for future read-only match operations."""
node = self.root_node
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_match_node = node
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
def _update_best_if_valid(node):
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
best_match_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
# HiCache: dead node (evicted + not backuped) — stop traversal
if child.evicted and not child.backuped:
break
prefix_len = child.key.match(key, page_size=self.page_size)
if prefix_len < len(child.key):
# Read-only: do not split, ignore partial match and stop
break
if not child.evicted:
value.append(child.component_data[BASE_COMPONENT_TYPE].value)
node = child
_update_best_if_valid(node)
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_match_node, best_value_len
def _match_prefix_helper(
self, key: RadixKey
) -> tuple[list[torch.Tensor], UnifiedTreeNode, int]:
) -> tuple[list[torch.Tensor], UnifiedTreeNode, UnifiedTreeNode, int]:
# Non-HiCache mode has only device-resident matches, so the scheduler
# device anchor follows the best match. In HiCache mode, host-backed
# nodes can also match, so we separately track the best device-resident
# match for scheduler prefix indices and locking.
node = self.root_node
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_match_node = node
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
best_match_device_node = node
best_match_device_value_len = 0
separate_device_match = self.cache_controller is not None
if separate_device_match:
validators = tuple(
comp.create_match_validator() for comp in self._components_tuple
)
device_validators = tuple(
comp.create_match_validator(match_device_only=True)
for comp in self._components_tuple
)
else:
validators = tuple(
comp.create_match_validator(match_device_only=True)
for comp in self._components_tuple
)
def _all_valid(validators, node):
return all([v(node) for v in validators])
def _update_best_if_valid(node):
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
nonlocal best_match_node
nonlocal best_match_device_value_len, best_match_device_node
matched = _all_valid(validators, node)
if matched:
best_match_node = node
if not separate_device_match:
if matched:
best_match_device_value_len = len(value)
best_match_device_node = node
return
if _all_valid(device_validators, node):
best_match_device_value_len = len(value)
best_match_device_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
@@ -668,14 +663,21 @@ class UnifiedRadixCache(BasePrefixCache):
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_match_node, best_value_len
return (
value,
best_match_node,
best_match_device_node,
best_match_device_value_len,
)
def _match_post_processor(
self,
params: MatchPrefixParams,
value: list[torch.Tensor],
best_match_node: UnifiedTreeNode,
best_value_len: int,
best_match_device_node: UnifiedTreeNode,
best_match_device_value_len: int,
) -> MatchResult:
node_update = best_match_node
for comp in self._components_tuple:
@@ -691,23 +693,21 @@ class UnifiedRadixCache(BasePrefixCache):
cur_time -= 0.00001
node_update = node_update.parent
# Walk up to find last_device_node
last_device_node = best_match_node
while last_device_node is not self.root_node and last_device_node.evicted:
last_device_node = last_device_node.parent
# Walk up to find last_host_node for full component.
if self.cache_controller is None:
last_host_node = best_match_device_node
else:
last_host_node = best_match_node
while last_host_node is not self.root_node and not last_host_node.backuped:
last_host_node = last_host_node.parent
# Walk up to find last_host_node
last_host_node = best_match_node
while last_host_node is not self.root_node and not last_host_node.backuped:
last_host_node = last_host_node.parent
if best_value_len > 0:
device_indices = torch.cat(value[:best_value_len])
if best_match_device_value_len > 0:
device_indices = torch.cat(value[:best_match_device_value_len])
else:
device_indices = self._empty_match_result.device_indices
result = MatchResult(
device_indices=device_indices,
last_device_node=last_device_node,
last_device_node=best_match_device_node,
last_host_node=last_host_node,
best_match_node=best_match_node,
host_hit_length=0,
@@ -718,7 +718,7 @@ class UnifiedRadixCache(BasePrefixCache):
result=result,
params=params,
value_chunks=value,
best_value_len=best_value_len,
best_value_len=best_match_device_value_len,
)
return result
@@ -1219,10 +1219,10 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node: UnifiedTreeNode,
mem_quota: Optional[int] = None,
req=None,
) -> Optional[torch.Tensor]:
) -> bool:
"""Load evicted KV data from host back to device (H→D)."""
if self.cache_controller is None:
return None
return False
# Build KV transfer
kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
@@ -1255,7 +1255,7 @@ class UnifiedRadixCache(BasePrefixCache):
mem_quota is not None and kv_tokens > mem_quota + result.delta
):
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
return False
avail = self.token_to_kv_pool_allocator.available_size()
if avail < kv_tokens:
@@ -1263,7 +1263,7 @@ class UnifiedRadixCache(BasePrefixCache):
result = self.evict(EvictParams(num_tokens=needed))
if result.num_tokens_evicted < needed:
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
return False
# Load H→D
aux_xfers = [x for xfers in comp_xfers.values() for x in xfers]
@@ -1276,7 +1276,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.dec_lock_ref(best_match_node, ancestor_lock_params)
if device_indices is None:
return None
return False
# Commit: each component gets only its own transfers
kv_xfer.device_indices = device_indices
@@ -1297,7 +1297,7 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node,
self.inc_lock_ref(best_match_node).to_dec_params(),
)
return device_indices
return True
def _build_sidecar_transfers(
self,
@@ -1432,25 +1432,41 @@ class UnifiedRadixCache(BasePrefixCache):
best_match_node = params.best_match_node
mem_quota = params.mem_quota
req = params.req
assert req is not None
last_best_match_device_node = req.last_node
def _collect_new_prefix_indices() -> torch.Tensor:
prefix_chunks: list[torch.Tensor] = []
node = best_match_node
while node is not last_best_match_device_node:
value = node.component_data[BASE_COMPONENT_TYPE].value
assert value is not None
prefix_chunks.append(value)
node = node.parent
if not prefix_chunks:
return self._empty_match_result.device_indices
prefix_chunks.reverse()
return torch.cat(prefix_chunks)
if best_match_node.evicted or params.host_hit_length > 0:
loading_values = self.load_back(best_match_node, mem_quota, req=req)
if loading_values is not None:
if self.load_back(best_match_node, mem_quota, req=req):
new_indices = _collect_new_prefix_indices()
if new_indices.numel() == 0:
return (
self._empty_match_result.device_indices,
last_best_match_device_node,
)
logger.debug(
"init_load_back success: loaded %d tokens for node %d",
len(loading_values),
len(new_indices),
best_match_node.id,
)
return loading_values, best_match_node
# Fallback: walk up to non-evicted ancestor
# TODO(ispobock): The fallback path is not correct. The last_device_node should consider all the components.
while best_match_node is not self.root_node and best_match_node.evicted:
best_match_node = best_match_node.parent
return new_indices, best_match_node
return (
self._empty_match_result.device_indices,
best_match_node,
last_best_match_device_node,
)
def check_hicache_events(self) -> None: