chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Keep CP HiCache valid tails page-owned during extension

Browse Source 
CP shared KV HiCache now treats non-page-aligned valid-tail nodes as page-owned when a later request extends beyond them. The prefix probe, match path, insert path, and prepared backup start now agree on flooring the reusable prefix to the previous physical page boundary, so prepared backup metadata cannot start mid-page or fail to attach at insertion.

Duplicate frees under CP HiCache now go through a page-safe free helper. Insert and unfinished duplicate ranges free only fully unprotected pages; no-insert completion still releases the right tail owned only by the finishing request.

Constraint: CP HiCache allocator frees whole physical pages even when called with token-granular locs.

Rejected: Partial-page sharing/refcounting | too complex for the current page-as-minimum-unit contract.

Rejected: Fix only prepare_write_backup_for_req | match_prefix and insert would still expose exact valid-tail hits and desynchronize prepared backup length.

Confidence: medium

Scope-risk: moderate

Directive: Do not expose non-page-aligned CP valid-tail hits to extending requests unless partial-page ownership is explicitly implemented end-to-end.

Tested: remote g0034 py_compile for touched files

Tested: remote g0034 test_cp_hicache_metadata.py 97 passed

Tested: remote g0034 test_cp_shared_kv_runtime.py 73 passed

Not-tested: test_cp_shared_kv_layout.py aborts during installed sgl_kernel architecture-specific op loading before assertions
This commit is contained in:
laoyao0822
2026-05-29 21:49:17 +08:00
parent 2a9dfcca6f
commit 40cf691c78
4 changed files with 414 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+278
View File
@@ -1344,6 +1344,41 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 6)
self.assertEqual(cache.cache_controller.reservations[0][0].tolist(), list(range(6)))
def test_prepare_write_backup_for_req_floors_mid_page_prefix_hit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=180
)
]
)
cache._probe_existing_radix_prefix_len_no_split = lambda key: 6
req = types.SimpleNamespace(
rid="rid-tail-extend",
fill_ids=list(range(10)),
cache_protected_len=6,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
self.assertIsNotNone(req.cp_hicache_prepared_backup)
self.assertEqual(
cache.cache_controller.reservations[0][0].tolist(),
list(range(4, 10)),
)
def test_cache_finished_req_keeps_cp_valid_tail_insert_key(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable_finished_insert = False
@@ -1382,6 +1417,106 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertEqual(inserted[0].value.tolist(), list(range(6)))
self.assertEqual([indices.tolist() for indices in freed], [[], []])
def test_cache_finished_req_cp_insert_duplicate_free_skips_partial_page(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable_finished_insert = False
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
)
allocator = RecordingTokenAllocator()
cache.token_to_kv_pool_allocator = allocator
cache.insert = lambda params: types.SimpleNamespace(prefix_len=3)
cache.dec_lock_ref = lambda node: None
req = types.SimpleNamespace(
origin_input_ids=list(range(6)),
output_ids=[],
req_pool_idx=0,
extra_key=None,
cache_protected_len=1,
last_node=TreeNode(),
cp_hicache_prepared_backup=None,
pop_committed_kv_cache=lambda: 6,
)
cache.cache_finished_req(req)
self.assertEqual(allocator.freed, [])
def test_cache_finished_req_cp_no_insert_frees_only_full_unprotected_pages(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable_finished_insert = False
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(12, dtype=torch.int64).view(1, 12)
)
allocator = RecordingTokenAllocator()
cache.token_to_kv_pool_allocator = allocator
cache.dec_lock_ref = lambda node: None
req = types.SimpleNamespace(
origin_input_ids=list(range(9)),
output_ids=[],
req_pool_idx=0,
extra_key=None,
cache_protected_len=1,
last_node=TreeNode(),
cp_hicache_prepared_backup=None,
pop_committed_kv_cache=lambda: 9,
)
cache.cache_finished_req(req, is_insert=False)
self.assertEqual(allocator.freed, [[4, 5, 6, 7, 8]])
def test_cache_unfinished_req_cp_duplicate_free_skips_partial_page(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
writes = []
class Pool:
req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8)
def write(self, index, values):
writes.append((index, values.clone()))
cache.req_to_token_pool = Pool()
allocator = RecordingTokenAllocator()
cache.token_to_kv_pool_allocator = allocator
cache.insert = lambda params: types.SimpleNamespace(prefix_len=3)
new_last_node = TreeNode()
cache.match_prefix = lambda params: types.SimpleNamespace(
device_indices=torch.arange(6, dtype=torch.int64),
last_device_node=new_last_node,
)
cache.dec_lock_ref = lambda node: None
cache.inc_lock_ref = lambda node: None
req = types.SimpleNamespace(
fill_ids=list(range(6)),
req_pool_idx=0,
extra_key=None,
cache_protected_len=1,
last_node=TreeNode(),
prefix_indices=torch.empty((0,), dtype=torch.int64),
cp_hicache_prepared_backup=None,
)
cache.cache_unfinished_req(req)
self.assertEqual(allocator.freed, [])
self.assertEqual(req.cache_protected_len, 6)
def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
@@ -2601,6 +2736,83 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(prefix_len, 4)
def test_cp_prepare_probe_floors_exact_valid_tail_when_request_extends(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.disable = False
cache.page_size = 4
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
root = TreeNode()
root.key = RadixKey([])
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 145
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_len = 6
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=6,
padded_len=8,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
root.children[(0, 1, 2, 3)] = node
prefix_len = cache._probe_existing_radix_prefix_len_no_split(
RadixKey(list(range(10)))
)
self.assertEqual(prefix_len, 4)
def test_cp_match_prefix_floors_exact_valid_tail_when_request_extends(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 4
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
cache.maybe_bigram_convert = lambda key: (key, None)
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.host_len = 0
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 146
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_value = None
node.host_len = 0
node.cp_hicache = None
root.children[(0, 1, 2, 3)] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(10)))))
self.assertEqual(result.device_indices.tolist(), [0, 1, 2, 3])
self.assertEqual(result.host_hit_length, 0)
self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3])
self.assertEqual(result.last_device_node.children[(4, 5)].key.token_ids, [4, 5])
def test_cp_insert_floors_backed_tail_split_to_page_boundary(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
@@ -2665,6 +2877,72 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertEqual(new_tail.key.token_ids, [4])
self.assertEqual(new_tail.value.tolist(), [4])
def test_cp_insert_extends_from_page_boundary_after_exact_valid_tail(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.disable = False
cache.is_eagle = False
cache.page_size = 4
cache.pending_host_backups = {}
cache.ongoing_write_through = {}
cache.cache_controller = types.SimpleNamespace(
has_draft_hicache=False,
write_policy="write_through",
)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache._inc_hit_count = lambda *args, **kwargs: None
cache._record_store_event = lambda node: None
cache.evictable_size_ = 0
cache.protected_size_ = 0
cache.enable_storage = False
cache.enable_kv_cache_events = False
cache.inc_node_lock_ref = lambda node: None
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.children = {}
cache.root_node = root
node = TreeNode()
node.id = 147
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = torch.arange(6, dtype=torch.int64)
node.host_len = 0
node.cp_hicache = None
root.children[(0, 1, 2, 3)] = node
reservation = make_write_reservation(
torch.arange(4, 10, dtype=torch.int64), node_id=148, host_start=190
)
prepared = PreparedCpHiCacheBackup(
node_id=148,
reservation=reservation,
metadata=reservation.metadata,
logical_len=6,
)
result = cache.insert(
InsertParams(
key=RadixKey(list(range(10))),
value=torch.arange(10, dtype=torch.int64),
cp_hicache_prepared_backup=prepared,
)
)
self.assertEqual(result.prefix_len, 4)
parent = root.children[(0, 1, 2, 3)]
self.assertEqual(parent.key.token_ids, [0, 1, 2, 3])
self.assertEqual(parent.children[(4, 5)].key.token_ids, [4, 5])
new_tail = parent.children[(4, 5, 6, 7)]
self.assertEqual(new_tail.key.token_ids, [4, 5, 6, 7, 8, 9])
self.assertTrue(prepared.attached)
self.assertIs(cache.pending_host_backups[148].node, new_tail)
def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = False