fix: preserve HiCache load-back allocator headroom
This commit is contained in:
@@ -442,6 +442,40 @@ class TestPrefillAdder(CustomTestCase):
|
||||
self.assertEqual(adder2.rem_chunk_tokens, 0) # 3 - 3 = 0
|
||||
self.assertEqual(result3, AddReqResult.OTHER)
|
||||
|
||||
def test_host_load_back_passes_mem_quota(self):
|
||||
running_batch = self.create_running_batch()
|
||||
self.mock_token_allocator.available_size.return_value = 512
|
||||
self.mock_tree_cache.init_load_back.return_value = (
|
||||
__import__("torch").tensor([1, 2, 3, 4], dtype=__import__("torch").int64),
|
||||
"loaded_node",
|
||||
)
|
||||
adder = self.create_adder(
|
||||
running_batch,
|
||||
page_size=64,
|
||||
rem_input_tokens=4096,
|
||||
rem_total_tokens=4096,
|
||||
)
|
||||
req = self.create_mock_req("req", priority=0, max_new_tokens=16)
|
||||
req.extend_input_len = 256
|
||||
req.host_hit_length = 128
|
||||
req.prefix_indices = __import__("torch").empty(
|
||||
(0,), dtype=__import__("torch").int64
|
||||
)
|
||||
req.last_node = object()
|
||||
req.last_host_node = object()
|
||||
req.fill_ids = list(range(256))
|
||||
req.cache_protected_len = 0
|
||||
req.set_extend_input_len = lambda value: setattr(req, "extend_input_len", value)
|
||||
req.sampling_params.ignore_eos = False
|
||||
|
||||
result = adder.add_one_req(
|
||||
req, has_chunked_req=False, truncation_align_size=None
|
||||
)
|
||||
|
||||
self.assertNotEqual(result, AddReqResult.NO_TOKEN)
|
||||
params = self.mock_tree_cache.init_load_back.call_args.args[0]
|
||||
self.assertEqual(params.mem_quota, 256)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -6,6 +6,7 @@ import torch
|
||||
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
||||
|
||||
@@ -52,7 +53,7 @@ class TestCpSharedKVLayout(unittest.TestCase):
|
||||
self.assertEqual(positions.tolist(), [11, 15])
|
||||
|
||||
|
||||
class TestCPSharedPagedAllocator(unittest.TestCase):
|
||||
class TestCPSharedPagedAllocator(CustomTestCase):
|
||||
def test_compute_owner_alloc_fallback_logs_every_event(self):
|
||||
from sglang.srt.mem_cache import common
|
||||
|
||||
@@ -95,6 +96,31 @@ class TestCPSharedPagedAllocator(unittest.TestCase):
|
||||
allocator.free(locs)
|
||||
self.assertEqual(allocator.available_size(), 64 * 8)
|
||||
|
||||
def test_allocator_state_str_reports_free_and_release_pages(self):
|
||||
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
|
||||
|
||||
allocator = CPSharedPagedTokenToKVPoolAllocator(
|
||||
logical_size=64 * 8,
|
||||
physical_size=64 * 2,
|
||||
page_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
device="cpu",
|
||||
kvcache=None,
|
||||
need_sort=True,
|
||||
cp_size=4,
|
||||
cp_rank=0,
|
||||
)
|
||||
allocator.free_pages = torch.tensor([1, 5], dtype=torch.int64)
|
||||
allocator.release_pages = torch.tensor([9], dtype=torch.int64)
|
||||
|
||||
state = allocator.allocator_state_str()
|
||||
|
||||
self.assertIn("allocator_available_size=192", state)
|
||||
self.assertIn("allocator_free_size=128", state)
|
||||
self.assertIn("free_pages=2", state)
|
||||
self.assertIn("allocator_release_size=64", state)
|
||||
self.assertIn("release_pages=1", state)
|
||||
|
||||
def test_compute_owner_page_assignment_matches_in_seq_zigzag(self):
|
||||
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
|
||||
build_in_seq_page_compute_owners,
|
||||
|
||||
@@ -39,6 +39,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
MatchPrefixParams,
|
||||
)
|
||||
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# Test constants
|
||||
DEFAULT_PAGE_SIZE = 4
|
||||
@@ -235,6 +236,16 @@ class TestTreeNode(unittest.TestCase):
|
||||
self.assertFalse(node2 < node1)
|
||||
|
||||
|
||||
class TestTreeNodeChildren(CustomTestCase):
|
||||
def test_missing_child_access_does_not_create_orphan_node(self):
|
||||
node = TreeNode()
|
||||
|
||||
with self.assertRaises(KeyError):
|
||||
_ = node.children["missing"]
|
||||
|
||||
self.assertEqual(node.children, {})
|
||||
|
||||
|
||||
class TestRadixCache(unittest.TestCase):
|
||||
"""Test cases for RadixCache class."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user