bugfix:fix unifiedcache c128 radix cache management (#39426)
This commit is contained in:
@@ -329,6 +329,14 @@ class C128SidecarComponent(TreeComponent):
|
||||
].clone()
|
||||
return cache_len + 1 if self.tree_core.is_eagle and cache_len > 0 else cache_len
|
||||
|
||||
def floor_cache_len(self, cache_len: int) -> int:
|
||||
logical_len = cache_len
|
||||
if self.tree_core.is_eagle and logical_len > 0:
|
||||
logical_len -= 1
|
||||
group_tokens = 128 * self.allocator.c128_attn_allocator.page_size
|
||||
floored = logical_len // group_tokens * group_tokens
|
||||
return floored + 1 if self.tree_core.is_eagle and floored > 0 else floored
|
||||
|
||||
def apply_component_action(self, action: ComponentAction) -> None:
|
||||
if isinstance(action, FreeComponentDeviceSlot):
|
||||
for page_ids in action.indices:
|
||||
|
||||
@@ -384,6 +384,11 @@ class TreeComponent(ABC):
|
||||
- Mamba: performs the copy-on-write into a per-request slot."""
|
||||
return result
|
||||
|
||||
def floor_cache_len(self, cache_len: int) -> int:
|
||||
"""Constrain the combined effective cache length after every
|
||||
component's `prepare_for_caching_req` truncation has been min'd."""
|
||||
return cache_len
|
||||
|
||||
def update_component_on_insert_overlap(
|
||||
self,
|
||||
node: UnifiedTreeNode,
|
||||
|
||||
@@ -995,6 +995,8 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
)
|
||||
if cl is not None:
|
||||
effective_cache_len = min(effective_cache_len, cl)
|
||||
for comp in self._components_tuple:
|
||||
effective_cache_len = comp.floor_cache_len(effective_cache_len)
|
||||
|
||||
# Truncate if needed; the tail free is deferred and batched with
|
||||
# the unaligned tail below so a shared boundary page is emitted once.
|
||||
@@ -1068,7 +1070,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
)
|
||||
ranges = [(free_from, len(kv_indices))]
|
||||
if tail_free_start is not None:
|
||||
ranges.append((tail_free_start, len(kv_indices_full)))
|
||||
if free_from < len(kv_indices) and tail_free_start <= len(kv_indices):
|
||||
# The two halves touch at the truncation boundary and share
|
||||
# that page; free the union as one range.
|
||||
ranges[0] = (free_from, len(kv_indices_full))
|
||||
else:
|
||||
ranges.append((tail_free_start, len(kv_indices_full)))
|
||||
self.free_kv_row(req.kv, ranges)
|
||||
else:
|
||||
self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
|
||||
@@ -1131,6 +1138,9 @@ class UnifiedRadixCache(BasePrefixCache):
|
||||
if cl is not None:
|
||||
effective_cache_len = min(effective_cache_len, cl)
|
||||
|
||||
for comp in self._components_tuple:
|
||||
effective_cache_len = comp.floor_cache_len(effective_cache_len)
|
||||
|
||||
radix_key = RadixKey(
|
||||
token_ids[:effective_cache_len],
|
||||
req.extra_key,
|
||||
|
||||
Reference in New Issue
Block a user