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/runtime/executor.py

@@ -7,6 +7,7 @@ from typing import List, Optional
 import torch
 from torch import Tensor
 
+from astrai.config.inference_config import InferenceConfig
 from astrai.extension.backend.attention import (
     CudaBackend,
     get_backend,
@@ -19,6 +20,7 @@ from astrai.inference.workspace import InferenceWorkspace
 from astrai.model.automodel import AutoModel
 
 logger = logging.getLogger(__name__)
+_config = InferenceConfig()
 
 
 @contextmanager
@@ -114,7 +116,7 @@ def _warmup_cuda_graphs(
     # shapes on first call (F.linear is the dominant cost).  This also warms
     # up the CUDA context (driver init) and compiles the graph-capture trace
     # that follows.  Custom .so kernels do NOT need this — they are pre-built.
-    warmup_len = 64
+    warmup_len = _config.prefill_warmup_len
     tid = "_warmup_prefill"
     if task_cache.task_alloc(tid, list(range(warmup_len))):
         with (
@@ -312,9 +314,20 @@ class Executor:
     ):
         tasks = sorted(tasks, key=lambda t: t.task_id)
         batch_sz = len(tasks)
+
+        # Validate batch size bounds
+        if batch_sz > self._workspace.max_batch_size:
+            raise ValueError(
+                f"Batch size {batch_sz} exceeds max_batch_size "
+                f"{self._workspace.max_batch_size}"
+            )
+
         prompt_lens = [len(t.prompt_ids) for t in tasks]
+
+        # Validate inputs before any resource allocation
         if any(start_pos >= prompt_len for prompt_len in prompt_lens):
             raise ValueError("prefill start_pos must precede every prompt end")
+
         q_lens = [prompt_len - start_pos for prompt_len in prompt_lens]
 
         input_ids = torch.tensor(
@@ -380,10 +393,18 @@ class Executor:
             return []
 
         b = len(tasks)
+
+        # Validate batch size bounds
+        if b > self._workspace.max_batch_size:
+            raise ValueError(
+                f"Batch size {b} exceeds max_batch_size "
+                f"{self._workspace.max_batch_size}"
+            )
+
         ws = self._workspace
         task_ids = [t.task_id for t in tasks]
-        cur_positions = [t.next_pos for t in tasks]
         task_sig = tuple(task_ids)
+        cur_positions = [t.next_pos for t in tasks]
 
         # ---- pre-replay: update input buffers in-place ----
 
@@ -392,9 +413,16 @@ class Executor:
         # slots — fill input ids device-to-device.  inference_mode guards
         # the read because the source was produced under sampling's
         # inference-mode context.
+        #
+        # ``cache_valid`` checks the decode cache's own task signature:
+        # req-index signatures in the cache manager are recycled when freed
+        # slots are reallocated to new tasks, so a fresh batch whose prefill
+        # re-bind coincides with a stale signature would otherwise replay a
+        # previous generation's tokens into ``input_ids``.
+        task_sig_match = self.task_cache.last_task_signature_matches(task_ids)
         cached = self._decode_cache
-        sig_match = cached is not None and cached.task_sig == task_sig
-        if sig_match and cached.last_tokens is not None:
+        cache_valid = cached is not None and cached.task_sig == task_sig
+        if task_sig_match and cache_valid and cached.last_tokens is not None:
             with torch.inference_mode():
                 input_ids = ws.fill_input_ids_from_device(cached.last_tokens)
         else:
@@ -404,9 +432,15 @@ class Executor:
 
         kv_cache = self.task_cache.bind(task_ids, ws)
 
-        reuse_decode_state = self.task_cache.bind_was_steady and sig_match
+        # Reuse sampling state only if all conditions hold:
+        # 1. KV bind detected steady increment (same req_indices, seq_lens +1)
+        # 2. Task signature matches (same task_ids in same order)
+        # 3. We have a valid cached decode state for THIS task set
+        reuse_decode_state = (
+            cache_valid and self.task_cache.bind_was_steady and task_sig_match
+        )
         if reuse_decode_state:
-            info = self._decode_cache.sampling_info
+            info = cached.sampling_info
             ws.position_ids[:b] += 1
         else:
             info = _build_sampling_batch_info(tasks, self.device)