From d14c02b0dc8d48aa70ca40ab09d78ba8ace04314 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 27 May 2026 02:41:45 +0800 Subject: [PATCH] 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 --- python/sglang/srt/managers/schedule_policy.py | 15 +++- .../managers/test_hicache_controller_cp.py | 90 ++++++++++++++----- .../unit/managers/test_prefill_adder.py | 76 +++++++++++++++- .../mem_cache/test_cp_hicache_metadata.py | 61 ++++++++++--- 4 files changed, 206 insertions(+), 36 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index f672ecab0..2abbbb667 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index d24f25847..7446038b7 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/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"): diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 4267e24fc..581eda883 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 99c1dfb85..3dbb0ae2d 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/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, [])