bugfix:fix unifiedcache c128 radix cache management (#39426)

This commit is contained in:
vstone-w
2026-09-16 16:47:42 +08:00
committed by GitHub
parent 8baeded6f3
commit 0e528dc9ff
3 changed files with 24 additions and 1 deletions
@@ -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,