refactor: harden inference cache state and attention dispatch

- split KVCache into phase-specific PrefillKVCache/DecodeKVCache types selected by start_pos
- unify steady-state detection in TaskCacheManager
- guard decode steady-state reuse with the cached task signature so recycled req slots cannot replay a prior generation's tokens and positions
- collapse attention backend fwd_decode/fwd_prefill into a single subclass-owned forward with a shared _check_fwd guard
- fix thread-safety gap in weight update and validate prefill inputs before KV allocation
- centralize magic constants in InferenceConfig and align docs with behavior
This commit is contained in:
2026-09-04 14:28:04 +08:00
parent e13fe53475
commit ae7fc3059a
12 changed files with 244 additions and 152 deletions
+5 -1
View File
@@ -7,6 +7,10 @@ from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, Deque, Dict, Generator, List, Literal, Optional
from astrai.config.inference_config import InferenceConfig
_config = InferenceConfig()
@dataclass
class TaskTiming:
@@ -123,7 +127,7 @@ class MetricsCollector:
stats = metrics.get_stats()
"""
def __init__(self, max_recent: int = 128):
def __init__(self, max_recent: int = _config.max_recent_tasks):
self._timings: Dict[str, TaskTiming] = {}
self._completed: Deque[TaskTiming] = deque(maxlen=max_recent)
self._lock = threading.Lock()