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

[UnifiedRadixCache] Fix HiCache load back start node (#25088)

Browse Source
This commit is contained in:
Ke Bao
2026-05-14 13:18:12 +08:00
committed by GitHub
parent c016246b0f
commit e4378ff37f
19 changed files with 268 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+4
View File
@@ -740,8 +740,10 @@ class Req(ReqDllmMixin):
self.extend_input_len = 0
# The relative logprob_start_len in an extend batch
self.extend_logprob_start_len = 0
# TODO(ispobock): rename to last_device_node
self.last_node: Any = None
self.last_host_node: Any = None
self.best_match_node: Any = None
self.host_hit_length = 0
# Tokens loaded from storage backend (L3) during prefetch for this request
self.storage_hit_length = 0
@@ -1041,12 +1043,14 @@ class Req(ReqDllmMixin):
self.prefix_indices,
self.last_node,
self.last_host_node,
self.best_match_node,
self.host_hit_length,
self.mamba_branching_seqlen,
) = (
match_result.device_indices,
match_result.last_device_node,
match_result.last_host_node,
match_result.best_match_node,
match_result.host_hit_length,
match_result.mamba_branching_seqlen,
)
python/sglang/srt/managers/schedule_policy.py
+3 -1
View File
@@ -105,11 +105,13 @@ def match_prefix_for_req(
req.prefix_indices,
req.last_node,
req.last_host_node,
req.best_match_node,
req.host_hit_length,
) = (
match_result.device_indices,
match_result.last_device_node,
match_result.last_host_node,
match_result.best_match_node,
match_result.host_hit_length,
)
if match_result.mamba_branching_seqlen is not None:
@@ -877,7 +879,7 @@ class PrefillAdder:
if req.host_hit_length > 0:
new_indices, req.last_node = self.tree_cache.init_load_back(
InitLoadBackParams(
last_host_node=req.last_host_node,
best_match_node=req.best_match_node,
host_hit_length=req.host_hit_length,
req=req,
)
python/sglang/srt/mem_cache/base_prefix_cache.py
+11 -2
View File
@@ -134,9 +134,9 @@ class DecLockRefResult:
@dataclasses.dataclass
class InitLoadBackParams:
"""Unified parameters for init_load_back across different cache types"""
"""Unified parameters for init_load_back across different cache types."""
last_host_node: Any
best_match_node: Any
host_hit_length: int
mem_quota: Optional[int] = None
req: Optional[Req] = None
@@ -151,6 +151,13 @@ class MatchResult(NamedTuple):
last_host_node : The last TreeNode on the host that was matched.
Note that if HiCache is not enabled,
this **must** be the same as `last_device_node`.
Reserved for L3 storage prefetch anchoring; L2 load_back
uses `best_match_node` instead.
best_match_node : Deepest node accepted by all component validators
during match_prefix. Anchor for every L2 host->device
load_back walk (FULL / SWA / ...). For legacy caches
that don't run multi-component validation, set this
equal to `last_host_node`.
host_hit_length : Length of the host cache hit. For pure-KV caches this is the
number of evicted KV tokens on CPU. For hybrid Mamba models this
is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
@@ -164,6 +171,7 @@ class MatchResult(NamedTuple):
device_indices: torch.Tensor
last_device_node: Any
last_host_node: Any
best_match_node: Any
host_hit_length: int = 0
mamba_branching_seqlen: Optional[int] = None
cache_protected_len: Optional[int] = None
@@ -180,6 +188,7 @@ def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult":
device_indices=match_result.device_indices[:0],
last_device_node=root,
last_host_node=root,
best_match_node=root,
host_hit_length=0,
)
python/sglang/srt/mem_cache/chunk_cache.py
+1
View File
@@ -69,6 +69,7 @@ class ChunkCache(BasePrefixCache):
device_indices=torch.empty((0,), dtype=torch.int64),
last_device_node=None,
last_host_node=None,
best_match_node=None,
)
def insert(self, params: InsertParams) -> InsertResult:
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+4 -1
View File
@@ -345,7 +345,7 @@ class HiMambaRadixCache(MambaRadixCache):
self,
params: InitLoadBackParams,
):
last_node = params.last_host_node
last_node = params.best_match_node
mem_quota = params.mem_quota
req = params.req
if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped):
@@ -932,6 +932,7 @@ class HiMambaRadixCache(MambaRadixCache):
device_indices=torch.empty((0,), dtype=torch.int64, device=self.device),
last_device_node=self.root_node,
last_host_node=self.root_node,
best_match_node=self.root_node,
host_hit_length=0,
)
@@ -1066,6 +1067,8 @@ class HiMambaRadixCache(MambaRadixCache):
device_indices=value,
last_device_node=last_device_node,
last_host_node=last_host_node,
# TODO(ispobock): use best_match_node as start node for load_back
best_match_node=last_host_node,
host_hit_length=host_hit_length,
mamba_branching_seqlen=mamba_branching_seqlen,
)
python/sglang/srt/mem_cache/hiradix_cache.py
+3 -1
View File
@@ -1046,7 +1046,7 @@ class HiRadixCache(RadixCache):
self,
params: InitLoadBackParams,
):
last_node = params.last_host_node
last_node = params.best_match_node
mem_quota = params.mem_quota
if last_node.evicted:
loading_values = self.load_back(last_node, mem_quota)
@@ -1251,6 +1251,8 @@ class HiRadixCache(RadixCache):
device_indices=value,
last_device_node=last_node,
last_host_node=last_host_node,
# TODO(ispobock): use best_match_node as start node for load_back
best_match_node=last_host_node,
host_hit_length=host_hit_length,
)
python/sglang/srt/mem_cache/mamba_radix_cache.py
+2
View File
@@ -488,6 +488,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
),
last_device_node=self.root_node,
last_host_node=self.root_node,
best_match_node=self.root_node,
)
value, last_node, best_value_len = self._match_prefix_helper(key)
@@ -1072,6 +1073,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
mamba_branching_seqlen=mamba_branching_seqlen,
)
python/sglang/srt/mem_cache/radix_cache.py
+2
View File
@@ -353,6 +353,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
),
last_device_node=self.root_node,
last_host_node=self.root_node,
best_match_node=self.root_node,
)
self._record_all_cleared_event()
@@ -413,6 +414,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
)
def insert(self, params: InsertParams) -> InsertResult:
python/sglang/srt/mem_cache/radix_cache_cpp.py
+1
View File
@@ -107,6 +107,7 @@ class RadixCacheCpp(BasePrefixCache):
device_indices=self._merge_tensor(device_indices_vec),
last_device_node=node_gpu,
last_host_node=node_cpu,
best_match_node=node_cpu,
host_hit_length=host_indices_length,
)
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
+1
View File
@@ -207,6 +207,7 @@ class LMCRadixCache(RadixCache):
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
)
return base_res
python/sglang/srt/mem_cache/swa_radix_cache.py
+2
View File
@@ -405,6 +405,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
),
last_device_node=self.root_node,
last_host_node=self.root_node,
best_match_node=self.root_node,
)
value, last_node, best_value_len = self._match_prefix_helper(key)
@@ -960,6 +961,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
device_indices=value,
last_device_node=last_node,
last_host_node=last_node,
best_match_node=last_node,
)
def _compact_single_child_chain(self, node: TreeNode) -> None:
python/sglang/srt/mem_cache/unified_cache_components/full_component.py
+21 -8
View File
@@ -152,12 +152,20 @@ class FullComponent(TreeComponent):
) -> IncLockRefResult:
ct = self.component_type
root = self.cache.root_node
delta = 0
cur = node
while cur != root:
cd = cur.component_data[ct]
assert cd.value is not None
# Skip the bottom evicted segment
while cur is not root and cur.component_data[ct].value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
cur = cur.parent
# Lock the device-on segment up to root
delta = 0
while cur is not root:
cd = cur.component_data[ct]
assert (
cd.value is not None
), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment"
if cd.lock_ref == 0:
key_len = len(cd.value)
self.cache.component_evictable_size_[ct] -= key_len
@@ -174,8 +182,12 @@ class FullComponent(TreeComponent):
) -> None:
ct = self.component_type
root = self.cache.root_node
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
cur = node
while cur != root:
if cur.id in skip_lock_node_ids:
cur = cur.parent
continue
cd = cur.component_data[ct]
assert cd.value is not None
assert cd.lock_ref > 0
@@ -203,15 +215,16 @@ class FullComponent(TreeComponent):
return None
if phase == CacheTransferPhase.LOAD_BACK:
# Walk evicted chain, collect host_values and nodes
# `node` is best_match_node. FULL device evict only from leaves,
# so once we hit a device-on node, everything above is also device-on
backed_up: list[torch.Tensor] = []
nodes: list = []
cur = node
while cur.evicted:
cd = cur.component_data[ct]
if cd.host_value is not None:
backed_up.append(cd.host_value)
nodes.append(cur)
assert cd.host_value is not None
backed_up.append(cd.host_value)
nodes.append(cur)
cur = cur.parent
backed_up.reverse()
nodes.reverse()
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+2 -3
View File
@@ -67,7 +67,7 @@ class MambaComponent(TreeComponent):
) -> MatchResult:
cow_mamba = params.cow_mamba
req = params.req
last_node = result.last_device_node
last_node = result.best_match_node
if len(value_chunks) > best_value_len:
chunk_size = get_global_server_args().mamba_cache_chunk_size
@@ -101,8 +101,7 @@ class MambaComponent(TreeComponent):
# HiCache: if mamba was evicted from device but has host backup,
# ensure host_hit_length >= 1 so load_back is triggered.
host_node = result.last_host_node
cd = host_node.component_data[self.component_type]
cd = last_node.component_data[self.component_type]
if cd.value is None and cd.host_value is not None:
result = result._replace(host_hit_length=max(result.host_hit_length, 1))
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+9 -5
View File
@@ -95,11 +95,12 @@ class SWAComponent(TreeComponent):
) -> MatchResult:
ct = self.component_type
n_swa = 0
node = result.last_host_node
node = result.best_match_node
root = self.cache.root_node
while node is not root and n_swa < self.sliding_window_size:
cd = node.component_data[ct]
if cd.value is None and cd.host_value is not None:
# TODO(ispobock): refactor host_hit_length usage
return result._replace(host_hit_length=max(result.host_hit_length, 1))
if cd.value is not None:
n_swa += len(cd.value)
@@ -440,11 +441,14 @@ class SWAComponent(TreeComponent):
]
if phase == CacheTransferPhase.LOAD_BACK:
# `node` is best_match_node; the SWA validator guarantees every
# ancestor within `sliding_window_size` has value or host_value.
n_swa = 0
backed_up: list[torch.Tensor] = []
nodes: list = []
while node is not self.cache.root_node and n_swa < self.sliding_window_size:
cd = node.component_data[ct]
cur = node
while cur is not self.cache.root_node and n_swa < self.sliding_window_size:
cd = cur.component_data[ct]
assert cd.host_value is not None or cd.value is not None
if cd.value is not None:
# device exists, skip it
@@ -452,9 +456,9 @@ class SWAComponent(TreeComponent):
else:
# host only, collect it
backed_up.append(cd.host_value)
nodes.append(node)
nodes.append(cur)
n_swa += len(cd.host_value)
node = node.parent
cur = cur.parent
if not backed_up:
return None
python/sglang/srt/mem_cache/unified_radix_cache.py
+41 -39
View File
@@ -290,6 +290,7 @@ class UnifiedRadixCache(BasePrefixCache):
),
last_device_node=self.root_node,
last_host_node=self.root_node,
best_match_node=self.root_node,
)
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
@@ -352,8 +353,10 @@ class UnifiedRadixCache(BasePrefixCache):
if len(key) == 0:
return self._empty_match_result
value, last_node, best_value_len = self._match_prefix_helper(key)
return self._match_post_processor(params, value, last_node, best_value_len)
value, best_match_node, best_value_len = self._match_prefix_helper(key)
return self._match_post_processor(
params, value, best_match_node, best_value_len
)
def insert(self, params: InsertParams) -> InsertResult:
if self.disable:
@@ -595,16 +598,16 @@ class UnifiedRadixCache(BasePrefixCache):
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_node = node
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_node
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
best_node = node
best_match_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
@@ -625,7 +628,7 @@ class UnifiedRadixCache(BasePrefixCache):
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_node, best_value_len
return value, best_match_node, best_value_len
def _match_prefix_helper(
self, key: RadixKey
@@ -634,16 +637,16 @@ class UnifiedRadixCache(BasePrefixCache):
child_key = key.child_key(self.page_size)
value: list[torch.Tensor] = []
best_value_len = 0
best_node = node
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_node
nonlocal best_value_len, best_match_node
if all(v(node) for v in validators):
best_value_len = len(value)
best_node = node
best_match_node = node
while len(key) > 0 and child_key in node.children:
child = node.children[child_key]
@@ -667,16 +670,16 @@ class UnifiedRadixCache(BasePrefixCache):
key = key[prefix_len:]
if len(key):
child_key = key.child_key(self.page_size)
return value, best_node, best_value_len
return value, best_match_node, best_value_len
def _match_post_processor(
self,
params: MatchPrefixParams,
value: list[torch.Tensor],
last_node: UnifiedTreeNode,
best_match_node: UnifiedTreeNode,
best_value_len: int,
) -> MatchResult:
node_update = last_node
node_update = best_match_node
for comp in self._components_tuple:
if comp.component_type == BASE_COMPONENT_TYPE:
continue # Full uses last_access_time, not LRU
@@ -691,12 +694,12 @@ class UnifiedRadixCache(BasePrefixCache):
node_update = node_update.parent
# Walk up to find last_device_node
last_device_node = last_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
last_host_node = last_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
@@ -708,6 +711,7 @@ class UnifiedRadixCache(BasePrefixCache):
device_indices=device_indices,
last_device_node=last_device_node,
last_host_node=last_host_node,
best_match_node=best_match_node,
host_hit_length=0,
)
@@ -1213,7 +1217,7 @@ class UnifiedRadixCache(BasePrefixCache):
def load_back(
self,
node: UnifiedTreeNode,
best_match_node: UnifiedTreeNode,
mem_quota: Optional[int] = None,
req=None,
) -> Optional[torch.Tensor]:
@@ -1222,15 +1226,12 @@ class UnifiedRadixCache(BasePrefixCache):
return None
# Build KV transfer
last_hit_node = node
kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
last_hit_node, CacheTransferPhase.LOAD_BACK
best_match_node, CacheTransferPhase.LOAD_BACK
)[0]
# Lock path & pre-evict if device pool is insufficient
nodes_to_load = kv_xfer.nodes_to_load
ancestor_node = nodes_to_load[0].parent if nodes_to_load else last_hit_node
result = self.inc_lock_ref(ancestor_node)
result = self.inc_lock_ref(best_match_node)
ancestor_lock_params = result.to_dec_params()
kv_tokens = len(kv_xfer.host_indices)
@@ -1240,7 +1241,7 @@ class UnifiedRadixCache(BasePrefixCache):
if comp.component_type == BASE_COMPONENT_TYPE:
continue
t = comp.build_hicache_transfers(
last_hit_node, CacheTransferPhase.LOAD_BACK, req=req
best_match_node, CacheTransferPhase.LOAD_BACK, req=req
)
if t:
comp_xfers[comp.component_type] = t
@@ -1255,7 +1256,7 @@ class UnifiedRadixCache(BasePrefixCache):
if (kv_tokens < self.load_back_threshold and not comp_xfers) or (
mem_quota is not None and kv_tokens > mem_quota + result.delta
):
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
avail = self.token_to_kv_pool_allocator.available_size()
@@ -1263,7 +1264,7 @@ class UnifiedRadixCache(BasePrefixCache):
needed = kv_tokens - avail
result = self.evict(EvictParams(num_tokens=needed))
if result.num_tokens_evicted < needed:
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
self.dec_lock_ref(best_match_node, ancestor_lock_params)
return None
# Load H→D
@@ -1271,32 +1272,32 @@ class UnifiedRadixCache(BasePrefixCache):
aux_xfers.extend(anchor_kv_shared_indices_xfers)
device_indices = self.cache_controller.load(
host_indices=kv_xfer.host_indices,
node_id=last_hit_node.id,
node_id=best_match_node.id,
extra_pools=aux_xfers or None,
)
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
self.dec_lock_ref(best_match_node, ancestor_lock_params)
if device_indices is None:
return None
# Commit: each component gets only its own transfers
kv_xfer.device_indices = device_indices
self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer(
last_hit_node,
best_match_node,
CacheTransferPhase.LOAD_BACK,
[kv_xfer],
)
for ct, xfers in comp_xfers.items():
self.components[ct].commit_hicache_transfer(
last_hit_node,
best_match_node,
CacheTransferPhase.LOAD_BACK,
xfers,
)
self._update_evictable_leaf_sets(ancestor_node)
self.ongoing_load_back[last_hit_node.id] = (
last_hit_node,
self.inc_lock_ref(last_hit_node).to_dec_params(),
self._update_evictable_leaf_sets(best_match_node)
self.ongoing_load_back[best_match_node.id] = (
best_match_node,
self.inc_lock_ref(best_match_node).to_dec_params(),
)
return device_indices
@@ -1384,27 +1385,28 @@ class UnifiedRadixCache(BasePrefixCache):
) -> tuple[torch.Tensor, UnifiedTreeNode]:
"""Prepare KV cache loading from host to device.
Returns (device_indices, last_node) tuple."""
last_node = params.last_host_node
best_match_node = params.best_match_node
mem_quota = params.mem_quota
req = params.req
if last_node.evicted or params.host_hit_length > 0:
loading_values = self.load_back(last_node, mem_quota, req=req)
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:
logger.debug(
"init_load_back success: loaded %d tokens for node %d",
len(loading_values),
last_node.id,
best_match_node.id,
)
return loading_values, last_node
return loading_values, best_match_node
# Fallback: walk up to non-evicted ancestor
while last_node is not self.root_node and last_node.evicted:
last_node = last_node.parent
# 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 (
self._empty_match_result.device_indices,
last_node,
best_match_node,
)
def check_hicache_events(self) -> None:
python/sglang/srt/session/streaming_session.py
+1
View File
@@ -259,6 +259,7 @@ class StreamingSession(BasePrefixCache):
device_indices=device_indices,
last_device_node=slot.virtual_node,
last_host_node=slot.virtual_node,
best_match_node=slot.virtual_node,
cache_protected_len=slot.cache_protected_len,
)