refactor: decouple task cache from PagePool and unify steady-state detection

- TaskCacheRegistry -> TaskCacheManager (independent, held by scheduler)
- TaskCacheState co-locates 5 parallel dicts into one dataclass
- AllocationStrategy base class + PagedStrategy subclass (page_size is a parameter)
- _rollback() helper for unified cleanup (no duplicate free paths)
- Task._kv_len + prefill_done property (explicit, no output_tokens proxy)
- Steady-state detection single-sourced in TaskCacheManager.bind()
- PagePool is now pure physical layer (no task knowledge)
- Removed dead _page_to_hash dict in RadixCache
This commit is contained in:
2026-08-08 22:43:06 +08:00
parent ca50fe4721
commit 3fa7e66676
10 changed files with 474 additions and 400 deletions
+17 -2
View File
@@ -77,8 +77,18 @@ class Task:
self.output_logprobs: List[float] = []
self.input_tokens: int = 0
self.output_tokens: int = 0
self._kv_len: int = 0
self._decoder: Optional[StreamDecoder] = None
def mark_prefill_done(self):
"""Prompt KV is materialized by prefill; first output sampled but
not yet written to KV."""
self._kv_len = self.input_tokens
def advance_kv(self):
"""One more position written to KV (after a decode forward)."""
self._kv_len += 1
def decode_new_token(self, tokenizer: AutoTokenizer) -> str:
"""Decode the last appended output token, buffering incomplete
multi-byte sequences across calls.
@@ -91,8 +101,13 @@ class Task:
@property
def next_pos(self) -> int:
# The first output is sampled from prefill and enters KV on the next step.
return self.input_tokens + max(0, len(self.output_ids) - 1)
"""KV position where the next decode step will write."""
return self._kv_len
@property
def prefill_done(self) -> bool:
"""True when all prompt KV entries are materialized."""
return self._kv_len >= self.input_tokens > 0
def is_finished(self, stop_ids: List[int]) -> bool:
if self.max_tokens is not None and self.output_tokens >= self.max_tokens: