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
+27 -16
View File
@@ -7,11 +7,20 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
import torch
from astrai.extension import ATTN_BACKEND, attn_backend
from astrai.inference.core.cache import PagePool
from astrai.inference.core.cache import PagePool, TaskCacheManager
from astrai.inference.core.workspace import InferenceWorkspace
from tests.extension.conftest import D, skip_no_kernel
def _mk_task_cache(pool: PagePool) -> TaskCacheManager:
return TaskCacheManager(
strategy=pool._strategy,
req_pool=pool._req_pool,
max_seq_len=pool.max_seq_len,
pool=pool,
)
def _ws(pool: PagePool) -> InferenceWorkspace:
return InferenceWorkspace(
pool.max_batch_size,
@@ -74,20 +83,21 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
dtype=torch.bfloat16,
)
task_cache = _mk_task_cache(cache)
ws = _ws(cache)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv1 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
task_cache.task_alloc("t1", prompt_ids[0])
task_cache.task_alloc("t2", prompt_ids[1])
kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode():
out_torch = model(
input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
)
cache.task_free("t1")
cache.task_free("t2")
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv2 = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
task_cache.task_free("t1")
task_cache.task_free("t2")
task_cache.task_alloc("t1", prompt_ids[0])
task_cache.task_alloc("t2", prompt_ids[1])
kv2 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(
@@ -138,10 +148,11 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
input_mask[i, : len(p)] = True
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
task_cache = _mk_task_cache(cache)
ws = _ws(cache)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv = cache.bind_tasks(["t1", "t2"], ws, start_pos=0)
task_cache.task_alloc("t1", prompt_ids[0])
task_cache.task_alloc("t2", prompt_ids[1])
kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode():
model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
@@ -151,15 +162,15 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
total_len = 9
dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
cache.task_extend("t1", 8)
cache.task_extend("t2", 6)
kv_t = cache.bind_tasks(["t1", "t2"], ws)
task_cache.task_extend("t1", 8)
task_cache.task_extend("t2", 6)
kv_t = task_cache.bind(["t1", "t2"], ws)
with torch.inference_mode():
out_torch = model(
dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
)
kv_c = cache.bind_tasks(["t1", "t2"], ws)
kv_c = task_cache.bind(["t1", "t2"], ws)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(