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

[UnifiedTree]: Fix SWA admission budget under-counts HiCache load-back consumption (#27391)

Browse Source
This commit is contained in:
Zhangheng
2026-06-07 10:47:08 +08:00
committed by GitHub
parent e57323cae9
commit fe548f36b0
9 changed files with 114 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+15
View File
@@ -816,7 +816,10 @@ class Req(ReqDllmMixin):
self.last_node: Any = None
self.last_host_node: Any = None
self.best_match_node: Any = None
# Per-component host hit lengths split off from host_hit_length:
self.host_hit_length = 0
self.swa_host_hit_length = 0
self.mamba_host_hit_length = 0
# Total cached prefix length (on-device prefix_indices + host_hit_length),
# capped at the max allowed prefix. Set during prefix matching at schedule
# time and used to estimate uncached tokens / sort by longest prefix for
@@ -1002,6 +1005,14 @@ class Req(ReqDllmMixin):
return self.output_ids[: self.finished_len]
return self.output_ids
def needs_host_load_back(self) -> bool:
"""Whether any cache layer has a host hit that needs L2 H2D load_back."""
return (
self.host_hit_length > 0
or self.swa_host_hit_length > 0
or self.mamba_host_hit_length > 0
)
def _cache_commit_len(self) -> int:
# Report only the prompt prefix so thinking + answer fall into the
# overallocated range and are reclaimed by release_kv_cache. #22373.
@@ -1107,6 +1118,8 @@ class Req(ReqDllmMixin):
self.last_host_node,
self.best_match_node,
self.host_hit_length,
self.swa_host_hit_length,
self.mamba_host_hit_length,
self.mamba_branching_seqlen,
) = (
match_result.device_indices,
@@ -1114,6 +1127,8 @@ class Req(ReqDllmMixin):
match_result.last_host_node,
match_result.best_match_node,
match_result.host_hit_length,
match_result.swa_host_hit_length,
match_result.mamba_host_hit_length,
match_result.mamba_branching_seqlen,
)
if match_result.cache_protected_len is not None:
python/sglang/srt/managers/schedule_policy.py
+18 -5
View File
@@ -108,12 +108,16 @@ def match_prefix_for_req(
req.last_host_node,
req.best_match_node,
req.host_hit_length,
req.swa_host_hit_length,
req.mamba_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,
match_result.swa_host_hit_length,
match_result.mamba_host_hit_length,
)
max_len = req._compute_max_prefix_len(len(token_ids))
req.num_matched_prefix_tokens = min(
@@ -558,7 +562,9 @@ class PrefillAdder:
return available_and_evictable - self.cur_rem_token_offset
def _swa_budget_for_req(self, extend_input_len: int) -> int:
def _swa_budget_for_req(
self, extend_input_len: int, swa_host_hit_length: int = 0
) -> int:
"""SWA pool budget per request. Only valid when is_hybrid_swa is True.
With chunked prefill + overlap scheduler, the peak SWA occupancy is:
@@ -573,7 +579,10 @@ class PrefillAdder:
alloc = min(extend_input_len, self.rem_chunk_tokens)
else:
alloc = extend_input_len
return max(alloc, self.tree_cache.sliding_window_size) + self.page_size
budget = max(alloc, self.tree_cache.sliding_window_size) + self.page_size
if swa_host_hit_length > 0:
budget += self.ceil_paged_tokens(swa_host_hit_length)
return budget
def ceil_paged_tokens(self, tokens: int) -> int:
return -(-tokens // self.page_size) * self.page_size
@@ -886,7 +895,9 @@ class PrefillAdder:
return AddReqResult.NO_TOKEN
if self.is_hybrid_swa:
swa_needed = self._swa_budget_for_req(req.extend_input_len)
swa_needed = self._swa_budget_for_req(
req.extend_input_len, swa_host_hit_length=req.swa_host_hit_length
)
if swa_needed >= self.rem_swa_tokens:
return AddReqResult.NO_TOKEN
@@ -906,11 +917,13 @@ class PrefillAdder:
return AddReqResult.NO_TOKEN
if self.is_hybrid_swa:
swa_needed = self._swa_budget_for_req(req.extend_input_len)
swa_needed = self._swa_budget_for_req(
req.extend_input_len, swa_host_hit_length=req.swa_host_hit_length
)
if swa_needed >= self.rem_swa_tokens:
return AddReqResult.NO_TOKEN
if req.host_hit_length > 0:
if req.needs_host_load_back():
new_indices, req.last_node = self.tree_cache.init_load_back(
InitLoadBackParams(
best_match_node=req.best_match_node,
python/sglang/srt/mem_cache/base_prefix_cache.py
+10 -5
View File
@@ -167,11 +167,12 @@ class MatchResult(NamedTuple):
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
host hit still triggers load-back without adding a separate field.
0 if HiCache is not enabled.
host_hit_length : Number of Full-KV tokens that hit on host (CPU) and need to be
loaded back to device. Pure-KV cache semantics;
swa_host_hit_length : Number of SWA tokens that hit on host (within the sliding
window) and will be load-back into the SWA device pool.
mamba_host_hit_length: Number of Mamba slots that hit on host and will be load-back
into the Mamba device pool. Typically 0 or 1.
mamba_branching_seqlen: The mamba radix cache branching point, which is the longest
page-aligned position that could've been cache hit if there
exists a mamba state.
@@ -182,6 +183,8 @@ class MatchResult(NamedTuple):
last_host_node: Any
best_match_node: Any
host_hit_length: int = 0
swa_host_hit_length: int = 0
mamba_host_hit_length: int = 0
mamba_branching_seqlen: Optional[int] = None
cache_protected_len: Optional[int] = None
@@ -199,6 +202,8 @@ def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult":
last_host_node=root,
best_match_node=root,
host_hit_length=0,
swa_host_hit_length=0,
mamba_host_hit_length=0,
)
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+2 -2
View File
@@ -1041,7 +1041,6 @@ class HiMambaRadixCache(MambaRadixCache):
mamba_host_hit = (
1 if (last_host_node.mamba_evicted and last_host_node.mamba_backuped) else 0
)
host_hit_length = max(kv_host_hit_length, mamba_host_hit)
mamba_node = best_last_node
if cow_mamba and mamba_node.mamba_value is not None:
@@ -1069,7 +1068,8 @@ class HiMambaRadixCache(MambaRadixCache):
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,
host_hit_length=kv_host_hit_length,
mamba_host_hit_length=mamba_host_hit,
mamba_branching_seqlen=mamba_branching_seqlen,
)
python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py
+4 -2
View File
@@ -113,10 +113,12 @@ class MambaComponent(TreeComponent):
req.mamba_needs_clear = False
# HiCache: if mamba was evicted from device but has host backup,
# ensure host_hit_length >= 1 so load_back is triggered.
# ensure mamba_host_hit_length >= 1 so load_back is triggered.
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))
result = result._replace(
mamba_host_hit_length=max(result.mamba_host_hit_length, 1)
)
return result._replace(mamba_branching_seqlen=branching_seqlen)
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+12 -3
View File
@@ -128,20 +128,29 @@ class SWAComponent(TreeComponent):
) -> MatchResult:
ct = self.component_type
n_swa = 0
swa_host_hit = 0
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)
elif cd.host_value is not None:
# TODO(hzh): load_back may currently restore a full host-tombstone
# segment whose length exceeds sliding_window_size. Once
# load_back is constrained to fetch only one sliding window
# worth of pages, cap swa_host_hit at sliding_window_size
# here so the scheduler budget matches the actual device-pool
# consumption.
swa_host_hit += len(cd.host_value)
n_swa += len(cd.host_value)
else:
break
node = node.parent
if swa_host_hit > 0:
return result._replace(
swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit)
)
return result
def update_component_on_insert_overlap(
python/sglang/srt/mem_cache/unified_radix_cache.py
+8 -1
View File
@@ -2330,7 +2330,14 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
prefix_chunks.reverse()
return torch.cat(prefix_chunks)
if best_match_node.evicted or params.host_hit_length > 0:
if (
best_match_node.evicted
or params.host_hit_length > 0
or (
req is not None
and (req.swa_host_hit_length > 0 or req.mamba_host_hit_length > 0)
)
):
if self.load_back(best_match_node, mem_quota, req=req):
new_indices = _collect_new_prefix_indices()
if new_indices.numel() == 0: