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:
@@ -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:
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user