[HiCache] fix: release write-through lock_ref during decode (#20049)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user