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
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

View File

@@ -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:

View File

@@ -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 = (

View File

@@ -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))

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