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
+23
View File
@@ -0,0 +1,23 @@
"""Inference engine configuration."""
from astrai.config.base import BaseConfig
class InferenceConfig(BaseConfig):
"""Configuration for inference workspace and execution parameters.
Centralizes magic constants previously scattered across inference modules.
Args:
max_splits (int): Maximum number of splits for split-KV attention (decode partial results). Defaults to 32.
q_tile_rows (int): Number of rows per Q tile in prefill ragged batching. Defaults to 64.
prefill_warmup_len (int): Prompt length for prefill warmup (cuBLAS auto-tuning). Defaults to 64.
default_rep_window (int): Default repetition penalty window size for frequency penalty. Defaults to 64.
max_recent_tasks (int): Maximum number of recent tasks tracked for aggregate statistics. Defaults to 128.
"""
max_splits: int = 32
q_tile_rows: int = 64
prefill_warmup_len: int = 64
default_rep_window: int = 64
max_recent_tasks: int = 128