feat: version rollout weight updates
Track a monotonic policy version across optimizer steps, scheduler updates, and rollout results. Serialize synchronous generation with weight acknowledgements and invalidate reusable prefix KV entries so cached samples remain attributable to the behavior policy that generated them.
This commit is contained in:
Vendored
+8
@@ -338,6 +338,14 @@ class TaskCacheManager:
|
||||
if state is not None:
|
||||
self._strategy.record_hashes(state, prompt_ids, start_logical_page)
|
||||
|
||||
def invalidate_cache(self) -> int:
|
||||
"""Drop reusable KV entries once all task-owned entries are released."""
|
||||
if self._states:
|
||||
raise RuntimeError("Cannot invalidate KV cache while tasks are active")
|
||||
self._bind_state = None
|
||||
self._bind_was_steady = False
|
||||
return self._strategy.invalidate_cache()
|
||||
|
||||
@staticmethod
|
||||
def task_cacheable_ids(task_id: str, prompt_ids: List[int], output_ids: List[int]):
|
||||
return list(prompt_ids) + list(output_ids[:-1])
|
||||
|
||||
Vendored
+22
@@ -89,6 +89,19 @@ class Allocator:
|
||||
if idx in self._lru:
|
||||
self._lru.move_to_end(idx)
|
||||
|
||||
def clear_cached(self) -> int:
|
||||
"""Release every unreferenced LRU page back to the free pool."""
|
||||
with self._lock:
|
||||
cached = list(self._lru)
|
||||
self._lru.clear()
|
||||
for idx in cached:
|
||||
if self._refs[idx] != 0:
|
||||
raise RuntimeError("Cannot invalidate a referenced cache page")
|
||||
if self.on_evict:
|
||||
self.on_evict(idx)
|
||||
self._free_mask |= 1 << idx
|
||||
return len(cached)
|
||||
|
||||
|
||||
class RadixNode:
|
||||
"""A page-aligned edge in the CPU-side prefix radix trie."""
|
||||
@@ -200,6 +213,10 @@ class AllocationStrategy(ABC):
|
||||
start: int,
|
||||
) -> None: ...
|
||||
|
||||
def invalidate_cache(self) -> int:
|
||||
"""Drop reusable KV entries after an inference weight update."""
|
||||
return 0
|
||||
|
||||
|
||||
class ContiguousStrategy(AllocationStrategy):
|
||||
"""Static contiguous allocation: slots are pre-assigned at pool init.
|
||||
@@ -316,3 +333,8 @@ class PagedStrategy(AllocationStrategy):
|
||||
full = len(prompt_ids) // self._page_size
|
||||
for i in range(start, min(full, len(state.pages))):
|
||||
self._prefix.record(state.pages[i], prompt_ids, i)
|
||||
|
||||
def invalidate_cache(self) -> int:
|
||||
if self._prefix is None:
|
||||
return 0
|
||||
return self._alloc.clear_cached()
|
||||
|
||||
Reference in New Issue
Block a user