[HiCache] fix: release write-through lock_ref during decode (#20049)

This commit is contained in:
shuwenn
2026-03-16 14:49:31 +08:00
committed by GitHub
parent 39336f5812
commit 42f18fe560
5 changed files with 22 additions and 0 deletions

View File

@@ -2362,6 +2362,11 @@ class Scheduler(
batch.batch_is_full = False
return batch
# Eagerly release lock_ref on completed write-through nodes so they
# become evictable, improving batch scheduling headroom.
if self.enable_hierarchical_cache:
self.tree_cache.flush_write_through_acks()
# Check if decode out of memory
if (kv_full_retract_flag := not batch.check_decode_mem()) or (
TEST_RETRACT and self.forward_ct % TEST_RETRACT_INTERVAL == 0

View File

@@ -229,6 +229,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
"""
raise NotImplementedError()
def flush_write_through_acks(self) -> None:
"""Release lock_ref on radix-tree nodes whose write-through has completed.
Lightweight operation that only processes finished write acks.
No-op for caches without hierarchical write-through support.
"""
pass
def check_hicache_events(self) -> Any:
"""
Check HiCache related activities to update radix tree and synchronize across TP workers if needed

View File

@@ -343,6 +343,9 @@ class HiMambaRadixCache(MambaRadixCache):
def ready_to_load_host_cache(self) -> int:
return self.cache_controller.start_loading()
def flush_write_through_acks(self) -> None:
self.writing_check()
def check_hicache_events(self):
self.writing_check()
self.loading_check()

View File

@@ -1094,6 +1094,9 @@ class HiRadixCache(RadixCache):
"""
return self.cache_controller.start_loading()
def flush_write_through_acks(self) -> None:
self.writing_check()
def check_hicache_events(self):
self.writing_check()
self.loading_check()

View File

@@ -325,6 +325,9 @@ class SessionAwareCache(BasePrefixCache):
def ready_to_load_host_cache(self):
return self.inner.ready_to_load_host_cache()
def flush_write_through_acks(self) -> None:
return self.inner.flush_write_through_acks()
def check_hicache_events(self):
return self.inner.check_hicache_events()