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
+3 -1
View File
@@ -18,6 +18,7 @@ from typing import (
from tokenizers.decoders import DecodeStream
from astrai.config.inference_config import InferenceConfig
from astrai.inference.metrics import MetricsCollector
from astrai.tokenize.tokenizer import AutoTokenizer
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
from astrai.extension import AttentionBackend
STOP = object()
_config = InferenceConfig()
@dataclass(frozen=True)
@@ -85,7 +87,7 @@ class Task:
top_p: float = 1.0,
top_k: int = 50,
frequency_penalty: float = 0.0,
rep_window: int = 64,
rep_window: int = _config.default_rep_window,
backend: Optional["AttentionBackend"] = None,
):
self.task_id = task_id