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

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:

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,

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
+        host_hit_length :   Number of Full-KV tokens that hit on host (CPU) and need to be
-                            number of evicted KV tokens on CPU. For hybrid Mamba models this
+                            loaded back to device. Pure-KV cache semantics;
-                            is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
+        swa_host_hit_length  :   Number of SWA tokens that hit on host (within the sliding
-                            host hit still triggers load-back without adding a separate field.
+                            window) and will be load-back into the SWA device pool.
-                            0 if HiCache is not enabled.
+        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

@@ -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

@@ -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

@@ -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

@@ -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:

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):

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
+        self._apply_match_to_req(req, match)
-        req.last_node = match.last_device_node
-        req.best_match_node = match.best_match_node
-        req.host_hit_length = match.host_hit_length
 
         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
+        self._apply_match_to_req(req, match)
-        req.last_node = match.last_device_node
-        req.best_match_node = match.best_match_node
-        req.host_hit_length = match.host_hit_length
 
         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
+        self._apply_match_to_req(req, match)
-        req.last_node = match.last_device_node
-        req.best_match_node = match.best_match_node
-        req.host_hit_length = match.host_hit_length
 
         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
+        """finalize_match_result accumulates host_value lengths of SWA tombstones
-        within the trailing window is tombstoned (cd.value is None,
+        within the trailing sliding window into ``swa_host_hit_length``. Out-of-window
-        cd.host_value is not None). Out-of-window tombstones and chains fully
+        tombstones and chains fully on device must leave ``swa_host_hit_length`` at 0.
-        on device must leave host_hit_length untouched.
+        ``host_hit_length`` is Full-KV only and is never written by SWA.
-
-        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.
         """
         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