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