Count evictable device cache when gating HiCache load-back

HiCache host hits can be skipped before load-back when the quota gate only counts immediately free KV allocator space. Under CP shared-KV pressure most reusable capacity may be represented as evictable radix-cache leaves, so the gate can incorrectly reject a host hit and leave prefill with cached-token zero despite host residency. Count device evictable cache in the quota estimate while leaving actual owner-lane allocation and eviction checks in the load path.

Constraint: CP HiCache load-back still has to respect owner-lane allocation and allocator eviction semantics.

Rejected: Force load-back regardless of quota | would bypass the scheduler pressure signal and increase OOM risk.

Rejected: Treat cache-hit zero as a transfer issue | logs showed host hits were found but skipped by quota before transfer.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove evictable cache from load-back capacity accounting without checking CP HiCache host-hit behavior under device pressure.

Tested: git diff --check

Tested: remote g0034 container pytest -q test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_alloc_pages_with_owners.py (90 passed, 3 warnings)

Not-tested: Full ETE GLM5 CP+HiCache+EAGLE pressure run after this quota change

Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/managers/schedule_policy.py

@@ -501,8 +501,19 @@ class PrefillAdder:
 
     def _get_available_device_tokens_for_load_back(self) -> int:
         if self.is_hybrid_swa:
-            return self.token_to_kv_pool_allocator.full_available_size()
-        return self.token_to_kv_pool_allocator.available_size()
+            return (
+                self.token_to_kv_pool_allocator.full_available_size()
+                + self.tree_cache.full_evictable_size()
+            )
+        if self.is_hybrid_ssm_cache:
+            return (
+                self.token_to_kv_pool_allocator.available_size()
+                + self.tree_cache.full_evictable_size()
+            )
+        return (
+            self.token_to_kv_pool_allocator.available_size()
+            + self.tree_cache.evictable_size()
+        )
 
     def _get_load_back_mem_quota(self, real_input_tokens: int) -> int:
         reserve_tokens = real_input_tokens + self.page_size

test/registered/unit/managers/test_hicache_controller_cp.py

@@ -91,6 +91,20 @@ if "sgl_kernel.kvcacheio" not in sys.modules:
         setattr(kvcacheio_stub, name, lambda *args, **kwargs: None)
     sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub
 
+_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+for _schema in (
+    "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+    "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+    "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+):
+    try:
+        _sgl_kernel_lib.define(_schema)
+    except RuntimeError as exc:
+        if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
+            raise
+
 from sglang.srt.managers.cache_controller import HiCacheController
 from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
@@ -150,9 +164,11 @@ class FakeAllocator:
     def __init__(self, alloc_result=None):
         self.alloc_result = alloc_result
         self.alloc_calls = []
+        self.owner_alloc_calls = []
         self.frees = []
         self.cp_size = 4
         self.cp_rank = 1
+        self.page_size = 4
         self.device_pool = FakeDevicePool()
 
     def get_kvcache(self):
@@ -164,6 +180,14 @@ class FakeAllocator:
             return None
         return self.alloc_result[:need_size].clone()
 
+    def alloc_pages_with_owners(self, page_owners):
+        owners = list(page_owners)
+        self.owner_alloc_calls.append(owners)
+        if self.alloc_result is None:
+            return None
+        need_size = len(owners) * self.page_size
+        return self.alloc_result[:need_size].clone()
+
     def free(self, indices):
         self.frees.append(indices.clone())
         return len(indices)
@@ -286,7 +310,7 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
         logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64)
 
         with self.assertRaisesRegex(
-            ValueError, "(host_indices|physical_device_indices).*whole pages"
+            ValueError, "_write_cp expects page-aligned device_indices"
         ):
             controller.write(logical_locs, node_id=21)
 
@@ -425,17 +449,40 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
             logical_len=16,
             owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
             host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
-            page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
-            page_size=1,
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
         )
 
         device_indices = controller.load_cp([node], node_id=11)
         controller.start_loading()
 
         self.assertEqual(device_indices.tolist(), list(range(64, 80)))
-        self.assertEqual(allocator.alloc_calls, [16])
+        self.assertEqual(allocator.alloc_calls, [])
+        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
 
+    def test_cp_load_frees_unexpected_owner_allocator_length(self):
+        host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
+        allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))
+        controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
+        node = TreeNode()
+        node.host_len = 16
+        node.cp_hicache = CpHiCacheNodeMetadata(
+            logical_len=16,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
+        )
+
+        with self.assertRaisesRegex(
+            RuntimeError, "alloc_pages_with_owners returned unexpected length"
+        ):
+            controller.load_cp([node], node_id=111)
+
+        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
+        self.assertEqual(allocator.frees[0].tolist(), list(range(64, 76)))
+
     def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self):
         host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         draft_host_pool = FakeHostPool(
@@ -457,14 +504,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
             owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
             host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
             draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
-            page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
-            page_size=1,
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
         )
 
         device_indices = controller.load_cp([node], node_id=14)
         controller.start_loading()
 
         self.assertEqual(device_indices.tolist(), list(range(64, 80)))
+        self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
         self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
         self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
         self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
@@ -481,14 +529,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
             logical_len=4,
             owned_positions=torch.empty((0,), dtype=torch.int64),
             host_indices=torch.empty((0,), dtype=torch.int64),
-            page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
-            page_size=1,
+            page_owners=torch.tensor([0], dtype=torch.int8),
+            page_size=4,
         )
 
         device_indices = controller.load_cp([node], node_id=12)
         controller.start_loading()
 
         self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
+        self.assertEqual(allocator.owner_alloc_calls, [[0]])
         self.assertEqual(host_pool.loads, [])
         self.assertEqual(len(controller.ack_load_queue), 1)
 
@@ -499,14 +548,15 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
         host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
         node = TreeNode()
         node.host_len = 16
-        node.cp_hicache = type(
-            "CpHiCacheMetadataStub",
-            (),
-            {
-                "owned_positions": torch.tensor([4, 5, 6, 7], dtype=torch.int64),
-                "host_indices": host_indices,
-            },
-        )()
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=16,
+            owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
+        )
+        metadata.host_indices = host_indices
+        node.cp_hicache = metadata
 
         controller.load_cp([node], node_id=13)
 
@@ -530,8 +580,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
             logical_len=16,
             owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
             host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
-            page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
-            page_size=1,
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
         )
 
         with self.assertRaisesRegex(
@@ -549,8 +599,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
             logical_len=16,
             owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
             host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
-            page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
-            page_size=1,
+            page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
+            page_size=4,
         )
 
         with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):

test/registered/unit/managers/test_prefill_adder.py

@@ -1,11 +1,66 @@
 import sys
+import types
 import unittest
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
-for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
-    if _mod not in sys.modules:
-        sys.modules[_mod] = MagicMock()
+import torch
+
+if "sgl_kernel" not in sys.modules:
+    sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
+sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
+sys.modules["sgl_kernel"].__path__ = []
+if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
+
+    def _sgl_kernel_getattr(name):
+        if name.startswith("__"):
+            raise AttributeError(name)
+        fn = lambda *args, **kwargs: None
+        setattr(sys.modules["sgl_kernel"], name, fn)
+        return fn
+
+    sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
+
+if "sgl_kernel.kvcacheio" not in sys.modules:
+    sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
+
+for _name in (
+    "sgl_per_token_group_quant_8bit",
+    "sgl_per_token_group_quant_fp8",
+    "sgl_per_token_quant_fp8",
+    "fp8_blockwise_scaled_mm",
+    "fp8_scaled_mm",
+    "silu_and_mul",
+):
+    if not hasattr(sys.modules["sgl_kernel"], _name):
+        setattr(sys.modules["sgl_kernel"], _name, lambda *args, **kwargs: None)
+
+if "sgl_kernel.quantization" not in sys.modules:
+    quantization_stub = types.ModuleType("sgl_kernel.quantization")
+    for _name in (
+        "ggml_dequantize",
+        "ggml_moe_a8",
+        "ggml_moe_a8_vec",
+        "ggml_moe_get_block_size",
+        "ggml_mul_mat_a8",
+        "ggml_mul_mat_vec_a8",
+    ):
+        setattr(quantization_stub, _name, lambda *args, **kwargs: None)
+    sys.modules["sgl_kernel.quantization"] = quantization_stub
+
+_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+for _schema in (
+    "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+    "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+    "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+):
+    try:
+        _sgl_kernel_lib.define(_schema)
+    except RuntimeError as exc:
+        if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
+            raise
 
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder
@@ -38,6 +93,7 @@ class TestPrefillAdder(CustomTestCase):
         tree_cache.full_evictable_size.return_value = full_evictable_size
         tree_cache.swa_evictable_size.return_value = swa_evictable_size
         tree_cache.evictable_size.return_value = evictable_size
+        tree_cache.supports_mamba.return_value = False
         tree_cache.disable = False
         tree_cache.inc_lock_ref.return_value = IncLockRefResult()
         tree_cache.dec_lock_ref.return_value = DecLockRefResult()
@@ -480,6 +536,20 @@ class TestPrefillAdder(CustomTestCase):
         params = self.mock_tree_cache.init_load_back.call_args.args[0]
         self.assertEqual(params.mem_quota, 320)
 
+    def test_load_back_mem_quota_counts_evictable_device_tokens(self):
+        self.mock_tree_cache = self.create_tree_cache(evictable_size=90000)
+        self.mock_token_allocator = self.create_token_allocator(available_size=1024)
+        adder = self.create_adder(
+            self.create_running_batch(),
+            page_size=64,
+            tree_cache=self.mock_tree_cache,
+            token_to_kv_pool_allocator=self.mock_token_allocator,
+        )
+
+        quota = adder._get_load_back_mem_quota(real_input_tokens=65536)
+
+        self.assertEqual(quota, 90000 + 1024 - 65536 - 64)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -76,6 +76,20 @@ except (ImportError, RuntimeError):
 if "sgl_kernel.kvcacheio" not in sys.modules:
     sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
 
+_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+for _schema in (
+    "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+    "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
+    "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
+    "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+):
+    try:
+        _sgl_kernel_lib.define(_schema)
+    except RuntimeError as exc:
+        if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
+            raise
+
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
 from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode
@@ -483,6 +497,34 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
 
         self.assertTrue(cache._node_backuped(node))
 
+    def test_single_node_write_lock_updates_device_evictable_leaf_set(self):
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable = False
+        cache.root_node = TreeNode()
+        cache.root_node.key = RadixKey([])
+        cache.evictable_leaves = set()
+        cache.evictable_size_ = 4
+        cache.protected_size_ = 0
+
+        node = TreeNode()
+        node.parent = cache.root_node
+        node.key = RadixKey([1, 2, 3, 4])
+        node.value = torch.arange(4, dtype=torch.int64)
+        cache.root_node.children[1] = node
+        cache.evictable_leaves.add(node)
+
+        cache.inc_node_lock_ref(node)
+
+        self.assertNotIn(node, cache.evictable_leaves)
+        self.assertEqual(cache.evictable_size(), 0)
+        self.assertEqual(cache.protected_size(), 4)
+
+        cache.dec_node_lock_ref(node)
+
+        self.assertIn(node, cache.evictable_leaves)
+        self.assertEqual(cache.evictable_size(), 4)
+        self.assertEqual(cache.protected_size(), 0)
+
     def test_inc_hit_count_does_not_rewrite_cp_backed_node(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache._uses_cp_hicache = True
@@ -867,18 +909,15 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
         cache.root_node.children[1] = node
         cache.evictable_host_leaves.add(node)
 
-        all_done_states = iter([False, True])
-        cache._cp_all_ranks_true = lambda done: next(all_done_states, True)
-        cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids]
-        cache._cp_filter_all_ranks_safe_node_ids = (
-            lambda node_ids, is_safe, **_kwargs: [
-                node_id
-                for node_id in node_ids
-                if is_safe(cache._cp_node_by_id(node_id))
-            ]
-        )
+        all_done_states = iter([0, 1])
 
-        physical_freed = cache._cp_evict_host_for_physical_slots(0)
+        def fake_all_reduce(done, op=None, group=None):
+            done.fill_(next(all_done_states, 1))
+
+        with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce):
+            physical_freed = cache._evict_host_for_physical_slots(
+                0, synchronize_across_ranks=True
+            )
 
         self.assertEqual(physical_freed, 0)
         self.assertEqual(freed, [])