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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user