From 40cf691c78858e63b6b8baf6afbb209daf271463 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 21:49:17 +0800 Subject: [PATCH] Keep CP HiCache valid tails page-owned during extension 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 68 +++++ python/sglang/srt/mem_cache/hiradix_cache.py | 41 ++- python/sglang/srt/mem_cache/radix_cache.py | 37 ++- .../mem_cache/test_cp_hicache_metadata.py | 278 ++++++++++++++++++ 4 files changed, 414 insertions(+), 10 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 0c0231807..95496c41a 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -1151,6 +1151,44 @@ Tests needed: - Add a CUDA/remote regression or source-level guard proving CUDA hot path cannot call `pad_token_locs_to_page_boundary()` with a non-page-start span. +Additional finding while implementing C16: + +- Flooring only `prepare_write_backup_for_req()` is insufficient. The same + exact-valid-tail extension must be treated consistently by radix + `match_prefix()` and final `insert()`. +- Example with `page_size=64`: an existing 100-token valid-tail node and a new + 120-token request must expose/reuse only the first 64 tokens. If + `match_prefix()` exposes 100 or `insert()` later reuses 100, the prepared + backup from 64..119 either starts mid-page in host accounting or becomes + unattached/mismatched at insert time. +- Therefore C16 is a three-site contract: prefix probe, scheduler-visible + prefix match, and insertion duplicate accounting all floor exact valid-tail + extensions to the previous physical page boundary. + +Implemented C16: + +- Added a shared exact-valid-tail extension floor helper in `HiRadixCache`. +- `_probe_existing_radix_prefix_len_no_split()` now floors an exact + non-page-aligned valid-tail hit when the request extends beyond that tail. +- `match_prefix()` now exposes only the floored page-boundary prefix in the same + situation and splits the radix node at that page boundary. +- `insert()` applies the same floor, stops after the page-boundary split, and + inserts the recomputed suffix from that boundary so prepared CP backup length + remains attachable. +- `prepare_write_backup_for_req()` still guards stale scheduler + `cache_protected_len` by flooring the raw backup start after first checking + that there is an actual suffix to backup. + +Completed C16 tests: + +```text +remote g0034 container: + test_prepare_write_backup_for_req_floors_mid_page_prefix_hit + test_cp_prepare_probe_floors_exact_valid_tail_when_request_extends + test_cp_match_prefix_floors_exact_valid_tail_when_request_extends + test_cp_insert_extends_from_page_boundary_after_exact_valid_tail +``` + ## C17: Duplicate/no-insert frees are still token-granular under CP pages Finding: @@ -1194,6 +1232,36 @@ Tests needed: - `cache_unfinished_req()` duplicate-free path with a non-page-aligned protected prefix. +Implemented C17: + +- Added page-floor/ceil helpers for CP free ranges. +- Routed `cache_finished_req()` duplicate frees, finished no-insert frees, and + `cache_unfinished_req()` duplicate frees through a CP page-safe range helper. +- Under CP HiCache, duplicate free ranges now skip partial left/right pages and + only pass full unprotected pages to the page allocator. The no-insert + finished path skips the partial left protected-prefix page but still includes + the finished request's right tail page because no radix node retains it. This + prevents a token-granular free from releasing a physical page that still + contains protected prefix or radix-owned valid-tail KV without leaking + no-insert tail pages. + +Completed C17 tests: + +```text +remote g0034 container: + test_cache_finished_req_cp_insert_duplicate_free_skips_partial_page + test_cache_finished_req_cp_no_insert_frees_only_full_unprotected_pages + test_cache_unfinished_req_cp_duplicate_free_skips_partial_page +``` + +Not-tested C16/C17: + +- `test_cp_shared_kv_runtime.py` passed in the remote container. +- `test_cp_shared_kv_layout.py` was attempted in the same remote container, but + collection aborted inside installed `sgl_kernel` architecture-specific op + loading before any assertion ran. The CP HiCache metadata/regression suite + passed in the same container. + ## C18: `cache_protected_len` now mixes valid-token and physical-page semantics Finding: diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index fdb4f4518..46328d8b4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -48,6 +48,7 @@ from sglang.srt.mem_cache.radix_cache import ( TreeNode, ceil_to_page_len, compute_node_hash_values, + floor_to_page_len, split_node_hash_value, ) from sglang.srt.mem_cache.utils import convert_to_bigram_key @@ -2008,6 +2009,9 @@ class HiRadixCache(RadixCache): while len(key) > 0 and child_key in node.children: child = node.children[child_key] prefix_len = self.key_match_fn(child.key, key) + prefix_len = self._cp_floor_exact_valid_tail_extension_len( + child, prefix_len, len(key) + ) if prefix_len <= 0: break if prefix_len < len(child.key): @@ -2040,6 +2044,20 @@ class HiRadixCache(RadixCache): return prefix_len return prefix_len // self.page_size * self.page_size + def _cp_floor_exact_valid_tail_extension_len( + self, child: TreeNode, prefix_len: int, request_len: int + ) -> int: + if ( + not self._uses_cp_hicache + or self.page_size <= 1 + or prefix_len <= 0 + or prefix_len != len(child.key) + or prefix_len >= request_len + or prefix_len % self.page_size == 0 + ): + return prefix_len + return floor_to_page_len(prefix_len, self.page_size) + def prepare_write_backup_for_req(self, req) -> None: if self.disable or not self._uses_cp_hicache: return @@ -2077,9 +2095,10 @@ class HiRadixCache(RadixCache): is_bigram=self.is_eagle, ) ) - start = min(max(req.cache_protected_len, existing_prefix_len), len(keys)) - if len(keys) <= start: + raw_start = min(max(req.cache_protected_len, existing_prefix_len), len(keys)) + if len(keys) <= raw_start: return + start = floor_to_page_len(raw_start, self.page_size) kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, start : len(keys) @@ -3428,6 +3447,11 @@ class HiRadixCache(RadixCache): if self._is_pinned(child): child.pin_expiry = time.monotonic() + child.pin_ttl prefix_len = self.key_match_fn(child.key, key) + prefix_len = self._cp_floor_exact_valid_tail_extension_len( + child, prefix_len, len(key) + ) + if prefix_len <= 0: + break if prefix_len < len(child.key): if ( self._uses_cp_hicache @@ -3528,10 +3552,19 @@ class HiRadixCache(RadixCache): total_prefix_length = 0 while len(key) > 0 and child_key in node.children.keys(): + parent_node = node node = node.children[child_key] node.last_access_time = time.monotonic() node.priority = max(node.priority, priority) - prefix_len = self.key_match_fn(node.key, key) + raw_prefix_len = self.key_match_fn(node.key, key) + prefix_len = self._cp_floor_exact_valid_tail_extension_len( + node, raw_prefix_len, len(key) + ) + stop_after_page_floor = prefix_len != raw_prefix_len + if prefix_len <= 0: + if stop_after_page_floor: + node = parent_node + break if prefix_len == len(node.key): if node.evicted: @@ -3575,6 +3608,8 @@ class HiRadixCache(RadixCache): if len(key): child_key = self.get_child_key_fn(key) + if stop_after_page_floor: + break if len(key): use_prepared_cp_backup = ( diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 554a90a39..66ab4da1a 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -123,6 +123,12 @@ def ceil_to_page_len(length: int, page_size: int) -> int: return ((length + page_size - 1) // page_size) * page_size +def floor_to_page_len(length: int, page_size: int) -> int: + if page_size <= 1: + return length + return length // page_size * page_size + + class TreeNode: counter = 0 @@ -472,6 +478,22 @@ class RadixCache(BasePrefixCache): prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked) return InsertResult(prefix_len=prefix_len) + def _free_kv_indices_range( + self, + kv_indices: torch.Tensor, + start: int, + end: int, + *, + include_partial_tail: bool = False, + ): + if getattr(self, "_uses_cp_hicache", False): + start = ceil_to_page_len(start, self.page_size) + if not include_partial_tail: + end = floor_to_page_len(end, self.page_size) + if end < start: + end = start + self.token_to_kv_pool_allocator.free(kv_indices[start:end]) + def cache_finished_req(self, req: Req, is_insert: bool = True): """Cache request when it finishes.""" # In deterministic mode, disable finished request insertion to radix cache @@ -518,16 +540,19 @@ class RadixCache(BasePrefixCache): self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss") new_prefix_len = result.prefix_len # Free the duplicates that were already in the tree - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : new_prefix_len] + self._free_kv_indices_range( + kv_indices, req.cache_protected_len, new_prefix_len ) else: if prepared_cp_backup is not None and hasattr( self, "_rollback_prepared_cp_backup" ): self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert") - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : len(keys)] + self._free_kv_indices_range( + kv_indices, + req.cache_protected_len, + len(keys), + include_partial_tail=True, ) if prepared_cp_backup is not None: req.cp_hicache_prepared_backup = None @@ -587,9 +612,7 @@ class RadixCache(BasePrefixCache): req.cp_hicache_prepared_backup = None new_prefix_len = result.prefix_len - self.token_to_kv_pool_allocator.free( - kv_indices[req.cache_protected_len : new_prefix_len] - ) + self._free_kv_indices_range(kv_indices, req.cache_protected_len, new_prefix_len) # The prefix indices could be updated, reuse it match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) 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 0d6a027ac..e7d992b08 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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