diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 182aeeb46..ac55e445c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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: diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 7e1a74b7a..588261d7c 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 62a0252b2..f0bbca6c9 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 47818e930..7b56634cc 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -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, ) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py index 5b1758bbd..c1e51f96b 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/mamba_component.py @@ -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) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 1f2a68f40..228f65a6c 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -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( diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 75f8b6d4f..f63254906 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index f532b630c..f9410f0c9 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -83,6 +83,7 @@ class TestPrefillAdder(CustomTestCase): req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False req.finished.return_value = False + req.needs_host_load_back.return_value = False return req def create_adder(self, running_batch, **kwargs): diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 4c35c0174..232b11f46 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -546,6 +546,15 @@ class UnifiedRadixCacheSuite: req_to_token_pool.alloc([req]) return req + def _apply_match_to_req(self, req, match): + req.prefix_indices = match.device_indices + req.last_node = match.last_device_node + req.last_host_node = match.last_host_node + req.best_match_node = match.best_match_node + req.host_hit_length = match.host_hit_length + req.swa_host_hit_length = match.swa_host_hit_length + req.mamba_host_hit_length = match.mamba_host_hit_length + def _make_seq(self, start: int, num_pages: int) -> list[int]: """Page-aligned token sequence of num_pages pages.""" page_size = self.cfg.page_size @@ -2857,7 +2866,26 @@ class UnifiedRadixCacheSuite: self.assertIs(result.best_match_node, leaf) self.assertIs(result.last_device_node, parent) self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key)) - self.assertEqual(result.host_hit_length, 1) + self.assertEqual(result.host_hit_length, len(leaf.key)) + self.assertEqual(result.swa_host_hit_length, len(leaf.key)) + + tree, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain too short") + leaf = chain[-1] + parent = chain[-2] + tokens = self._match_tokens_for_chain(chain) + + self._set_aux_host_tombstone(tree, leaf, ComponentType.SWA) + + result = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + + self.assertIs(result.best_match_node, leaf) + self.assertIs(result.last_device_node, parent) + self.assertEqual(len(result.device_indices), len(tokens) - len(leaf.key)) + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.swa_host_hit_length, len(leaf.key)) def test_mamba_branching_seqlen_disabled_under_hicache(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: @@ -2911,10 +2939,7 @@ class UnifiedRadixCacheSuite: match = tree.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req) ) - req.prefix_indices = match.device_indices - req.last_node = match.last_device_node - req.best_match_node = match.best_match_node - req.host_hit_length = match.host_hit_length + self._apply_match_to_req(req, match) new_indices, new_node = tree.init_load_back( InitLoadBackParams( @@ -2955,10 +2980,7 @@ class UnifiedRadixCacheSuite: match = tree.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req) ) - req.prefix_indices = match.device_indices - req.last_node = match.last_device_node - req.best_match_node = match.best_match_node - req.host_hit_length = match.host_hit_length + self._apply_match_to_req(req, match) new_indices, new_node = tree.init_load_back( InitLoadBackParams( @@ -2996,10 +3018,7 @@ class UnifiedRadixCacheSuite: match = tree.match_prefix( MatchPrefixParams(key=RadixKey(array("q", tokens)), req=req) ) - req.prefix_indices = match.device_indices - req.last_node = match.last_device_node - req.best_match_node = match.best_match_node - req.host_hit_length = match.host_hit_length + self._apply_match_to_req(req, match) new_indices, new_node = tree.init_load_back( InitLoadBackParams( @@ -3113,14 +3132,10 @@ class UnifiedRadixCacheSuite: return tree, allocator, req_to_token_pool, chain, window_pages def test_hicache_swa_finalize_match_result(self): - """finalize_match_result bumps host_hit_length to 1 iff some SWA node - within the trailing window is tombstoned (cd.value is None, - cd.host_value is not None). Out-of-window tombstones and chains fully - on device must leave host_hit_length untouched. - - Sentinel only — never the real SWA token count, since SWA load-back - does not grow req.prefix_indices and any non-zero value gets - subtracted from extend_input_len in schedule_policy. + """finalize_match_result accumulates host_value lengths of SWA tombstones + within the trailing sliding window into ``swa_host_hit_length``. Out-of-window + tombstones and chains fully on device must leave ``swa_host_hit_length`` at 0. + ``host_hit_length`` is Full-KV only and is never written by SWA. """ if not self.cfg.has_swa: self.skipTest("requires SWA") @@ -3129,11 +3144,12 @@ class UnifiedRadixCacheSuite: tree, _, _, chain, window_pages = self._swa_finalize_setup() leaf = chain[-1] + ps = self.cfg.page_size swa_comp = tree.components[ComponentType.SWA] cases = [ ("all_on_device", None, 0), - ("tombstone_in_window", chain[-window_pages], 1), + ("tombstone_in_window", chain[-window_pages], ps), ("tombstone_outside_window", chain[-(window_pages + 1)], 0), ] for name, victim, expected in cases: @@ -3163,7 +3179,8 @@ class UnifiedRadixCacheSuite: value_chunks=[], best_value_len=0, ) - self.assertEqual(result.host_hit_length, expected) + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.swa_host_hit_length, expected) def test_hicache_swa_commit_load_back_rebuilds_mapping(self): """LOAD_BACK commit must: @@ -3312,6 +3329,7 @@ class UnifiedRadixCacheSuite: def test_hicache_swa_finalize_anchored_on_best_match_node(self): tree, _, _, y, x, _ = self._swa_anchor_setup() swa_comp = tree.components[ComponentType.SWA] + ps = self.cfg.page_size base = MatchResult( device_indices=torch.empty((0,), dtype=torch.int64, device=tree.device), @@ -3326,7 +3344,9 @@ class UnifiedRadixCacheSuite: value_chunks=[], best_value_len=0, ) - self.assertEqual(result.host_hit_length, 1) + # SWA host hit goes to swa_host_hit_length; host_hit_length stays at 0. + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.swa_host_hit_length, ps) def test_hicache_swa_temp_lock_does_not_release_restored_tombstone(self): """A temporary scheduler lock that skipped a SWA tombstone must not