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:
@@ -417,7 +417,7 @@ cycle belong under `TYPE_CHECKING`.
|
||||
|
||||
`astrai/extension/backend/attention.py` provides the backend abstraction:
|
||||
|
||||
- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
|
||||
- **`AttentionBackend`** (ABC): single abstract `forward`; each subclass branches on `fwd` ("decode" / "prefill" / None) internally, `_check_fwd` guards unknown modes
|
||||
- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
|
||||
- **`FlashAttnBackend`**: Optional flash-attn dispatch via `flash_attn_varlen_func` over gathered flat K/V.
|
||||
- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
|
||||
|
||||
Reference in New Issue
Block a user