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

astrai/inference/cache/pool.py

@@ -16,7 +16,13 @@ from typing import Dict, List, Optional
 
 import torch
 
-from astrai.inference.cache.buffer import KVCache, KVStorage, ReqToTokenPool
+from astrai.inference.cache.buffer import (
+    DecodeKVCache,
+    KVCache,
+    KVStorage,
+    PrefillKVCache,
+    ReqToTokenPool,
+)
 from astrai.inference.cache.strategy import (
     AllocationStrategy,
     Allocator,
@@ -31,6 +37,8 @@ from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace
 # continues to work unchanged after the file split.
 __all__ = [
     "KVCache",
+    "PrefillKVCache",
+    "DecodeKVCache",
     "KVStorage",
     "ReqToTokenPool",
     "Allocator",
@@ -232,7 +240,20 @@ class PagePool:
             )
             q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles]
             q_tile_to_index = workspace.q_tile_to_index[:n_tiles]
-            decode_o_part = decode_ml_part = decode_out = None
+
+            return PrefillKVCache(
+                k_buffer=self._storage.k_buffer,
+                v_buffer=self._storage.v_buffer,
+                req_to_token=self._req_pool.req_to_token,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens_t,
+                max_len=max(seq_lens),
+                kv_indptr=kv_indptr,
+                out_cache_loc=out_cache_loc,
+                qo_indptr=qo_indptr,
+                q_tile_to_batch=q_tile_to_batch,
+                q_tile_to_index=q_tile_to_index,
+            )
         else:
             # ---- decode: out_cache_loc is a single column (last position) ----
             write_pos = seq_lens_t - 1
@@ -241,27 +262,24 @@ class PagePool:
             out_cache_loc = ocl_buf[:b].reshape(-1)
             workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
             qo_indptr = workspace.qo_indptr[: b + 1]
-            q_tile_to_batch = q_tile_to_index = None
             decode_o_part = getattr(workspace, "decode_o_part", None)
             decode_ml_part = getattr(workspace, "decode_ml_part", None)
             decode_out = getattr(workspace, "decode_out", None)
 
-        return KVCache(
-            k_buffer=self._storage.k_buffer,
-            v_buffer=self._storage.v_buffer,
-            req_to_token=self._req_pool.req_to_token,
-            req_pool_indices=req_pool_indices,
-            seq_lens=seq_lens_t,
-            out_cache_loc=out_cache_loc,
-            max_len=max(seq_lens),
-            kv_indptr=kv_indptr,
-            qo_indptr=qo_indptr,
-            q_tile_to_batch=q_tile_to_batch,
-            q_tile_to_index=q_tile_to_index,
-            decode_o_part=decode_o_part,
-            decode_ml_part=decode_ml_part,
-            decode_out=decode_out,
-        )
+            return DecodeKVCache(
+                k_buffer=self._storage.k_buffer,
+                v_buffer=self._storage.v_buffer,
+                req_to_token=self._req_pool.req_to_token,
+                req_pool_indices=req_pool_indices,
+                seq_lens=seq_lens_t,
+                max_len=max(seq_lens),
+                kv_indptr=kv_indptr,
+                out_cache_loc=out_cache_loc,
+                qo_indptr=qo_indptr,
+                decode_o_part=decode_o_part,
+                decode_ml_part=decode_ml_part,
+                decode_out=decode_out,
+            )
 
 
 class TaskCacheManager:
@@ -385,8 +403,25 @@ class TaskCacheManager:
 
     @property
     def bind_was_steady(self) -> bool:
+        """True if the last bind was a steady-state increment (same tasks, +1 seq_lens)."""
         return self._bind_was_steady
 
+    def last_task_signature_matches(self, task_ids: List[str]) -> bool:
+        """Check if task_ids match the previous bind's signature.
+
+        Used by Executor to detect steady-state decode for device-to-device
+        token copy optimization.
+        """
+        if self._bind_state is None:
+            return False
+        prev_sig = self._bind_state.sig
+        # sig is tuple of req_indices, need to map task_ids to req_indices
+        try:
+            current_sig = tuple(self._states[tid].req_idx for tid in task_ids)
+            return prev_sig == current_sig
+        except KeyError:
+            return False
+
     # -- internals --
 
     def _rollback(self, state: TaskCacheState, task_id: str):